├── src └── lightbeam │ ├── __init__.py │ ├── misc.py │ ├── zernike.py │ ├── geom.py │ ├── screen.py │ ├── LPmodes.py │ ├── mesh.py │ ├── PIAA.py │ ├── optics.py │ └── prop.py ├── .gitattributes ├── tutorial ├── amr_comp.png ├── convergence.npy ├── run_bpm_example.py ├── convergence.py └── config_example.py ├── pyproject.toml ├── setup.cfg ├── LICENSE ├── README.md └── .gitignore /src/lightbeam/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /.gitattributes: -------------------------------------------------------------------------------- 1 | *.ipynb linguist-detectable=false 2 | -------------------------------------------------------------------------------- /tutorial/amr_comp.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/jw-lin/lightbeam/HEAD/tutorial/amr_comp.png -------------------------------------------------------------------------------- /tutorial/convergence.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/jw-lin/lightbeam/HEAD/tutorial/convergence.npy -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [build-system] 2 | requires = [ 3 | "setuptools>=42", 4 | "wheel" 5 | ] 6 | build-backend = "setuptools.build_meta" -------------------------------------------------------------------------------- /setup.cfg: -------------------------------------------------------------------------------- 1 | [metadata] 2 | name = lightbeam 3 | version = 0.0.1 4 | author = Jonathan Lin 5 | author_email = jon880@g.ucla.edu 6 | description = A lightweight beam propagator. 7 | long_description = file: README.md 8 | long_description_content_type = text/markdown 9 | url = https://github.com/jw-lin/lightbeam 10 | classifiers = 11 | Programming Language :: Python :: 3 12 | License :: OSI Approved :: MIT License 13 | Operating System :: OS Independent 14 | 15 | [options] 16 | package_dir = 17 | = src 18 | packages = find: 19 | python_requires = >=3.6 20 | 21 | [options.packages.find] 22 | where = src -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2023 Jonathan Lin 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # lightbeam 2 | Simulate light through weakly guiding waveguides using the finite-differences beam propagation method on an adaptive grid. 3 | 4 | ## installation 5 | Use pip: 6 | 7 | ``` 8 | pip install git+https://github.com/jw-lin/lightbeam.git 9 | ``` 10 | 11 | Update: 12 | 13 | ``` 14 | pip install --force-reinstall git+https://github.com/jw-lin/lightbeam.git 15 | ``` 16 | 17 | Python dependencies: `numpy`,`scipy`,`matplotlib`,`numba`,`numexpr`,`jupyter` 18 | 19 | ## getting started 20 | Check out the Python notebook in the `tutorial` folder for a quickstart guide. Direct link. 21 | 22 | ## references 23 | J. Shibayama, K. Matsubara, M. Sekiguchi, J. Yamauchi and H. Nakano, "Efficient nonuniform schemes for paraxial and wide-angle finite-difference beam propagation methods," in Journal of Lightwave Technology, vol. 17, no. 4, pp. 677-683, April 1999, doi: 10.1109/50.754799. 24 | 25 | The Beam-Propagation Method. (2015). In Beam Propagation Method for Design of Optical Waveguide Devices (pp. 22–70). doi:10.1002/9781119083405.ch2 26 | 27 | ## acknowledgments 28 | NSF grants 2109231, 2109232, 2308360, 2308361 29 | -------------------------------------------------------------------------------- /tutorial/run_bpm_example.py: -------------------------------------------------------------------------------- 1 | ''' example script for running the beamprop code in prop.py''' 2 | import numpy as np 3 | from lightbeam.mesh import RectMesh3D 4 | from lightbeam.prop import Prop3D 5 | from lightbeam.misc import normalize,overlap_nonu,norm_nonu 6 | from lightbeam import LPmodes 7 | import matplotlib.pyplot as plt 8 | from config_example import * 9 | 10 | if __name__ == "__main__": 11 | 12 | # mesh initialization (required) 13 | mesh = RectMesh3D(xw0,yw0,zw,ds,dz,num_PML,xw_func,yw_func) 14 | xg,yg = mesh.xy.xg,mesh.xy.yg 15 | 16 | mesh.xy.max_iters = max_remesh_iters 17 | mesh.sigma_max = sig_max 18 | 19 | # propagator initialization (required) 20 | prop = Prop3D(wl0,mesh,optic,n0) 21 | 22 | print('launch field') 23 | plt.imshow(np.real(u0)) 24 | plt.show() 25 | 26 | # run the propagator (required) 27 | u,u0 = prop.prop2end(u0,monitor_func=monitor_func,xyslice=None,zslice=None,writeto=writeto,ref_val=ref_val,remesh_every=remesh_every,dynamic_n0=dynamic_n0,fplanewidth=fplanewidth) 28 | 29 | # compute power in output ports (optional) 30 | 31 | xg,yg = np.meshgrid(mesh.xy.xa,mesh.xy.ya,indexing='ij') 32 | 33 | w = mesh.xy.get_weights() 34 | 35 | xg0,yg0 = np.meshgrid(mesh.xy.xa0,mesh.xy.ya0,indexing='ij') 36 | w0 = mesh.xy.dx0*mesh.xy.dy0 37 | 38 | modes = [] 39 | for x,y in zip(xpos,ypos): 40 | mode = norm_nonu(LPmodes.lpfield(xg-x,yg-y,0,1,rcore/scale,wl0,ncore,nclad),w) 41 | modes.append(mode) 42 | 43 | SMFpower=0 44 | print("final field power decomposition:") 45 | for i in range(len(modes)): 46 | _p = np.power(overlap_nonu(u,modes[i],w),2) 47 | print("mode"+str(i)+": ", _p) 48 | SMFpower += _p 49 | 50 | print("total power in SMFs: ", SMFpower) 51 | 52 | # plotting (optional) 53 | print("final field dist:") 54 | plt.imshow(np.abs(u0)[num_PML:-num_PML,num_PML:-num_PML]) 55 | plt.show() 56 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | .gitattributes 2 | # Byte-compiled / optimized / DLL files 3 | __pycache__/ 4 | *.py[cod] 5 | *$py.class 6 | 7 | # C extensions 8 | *.so 9 | 10 | # Distribution / packaging 11 | .Python 12 | build/ 13 | develop-eggs/ 14 | dist/ 15 | downloads/ 16 | eggs/ 17 | .eggs/ 18 | lib/ 19 | lib64/ 20 | parts/ 21 | sdist/ 22 | var/ 23 | wheels/ 24 | pip-wheel-metadata/ 25 | share/python-wheels/ 26 | *.egg-info/ 27 | .installed.cfg 28 | *.egg 29 | MANIFEST 30 | 31 | # PyInstaller 32 | # Usually these files are written by a python script from a template 33 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 34 | *.manifest 35 | *.spec 36 | 37 | # Installer logs 38 | pip-log.txt 39 | pip-delete-this-directory.txt 40 | 41 | # Unit test / coverage reports 42 | htmlcov/ 43 | .tox/ 44 | .nox/ 45 | .coverage 46 | .coverage.* 47 | .cache 48 | nosetests.xml 49 | coverage.xml 50 | *.cover 51 | *.py,cover 52 | .hypothesis/ 53 | .pytest_cache/ 54 | 55 | # Translations 56 | *.mo 57 | *.pot 58 | 59 | # Django stuff: 60 | *.log 61 | local_settings.py 62 | db.sqlite3 63 | db.sqlite3-journal 64 | 65 | # Flask stuff: 66 | instance/ 67 | .webassets-cache 68 | 69 | # Scrapy stuff: 70 | .scrapy 71 | 72 | # Sphinx documentation 73 | docs/_build/ 74 | 75 | # PyBuilder 76 | target/ 77 | 78 | # Jupyter Notebook 79 | .ipynb_checkpoints 80 | 81 | # IPython 82 | profile_default/ 83 | ipython_config.py 84 | 85 | # pyenv 86 | .python-version 87 | 88 | # pipenv 89 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 90 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 91 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 92 | # install all needed dependencies. 93 | #Pipfile.lock 94 | 95 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 96 | __pypackages__/ 97 | 98 | # Celery stuff 99 | celerybeat-schedule 100 | celerybeat.pid 101 | 102 | # SageMath parsed files 103 | *.sage.py 104 | 105 | # Environments 106 | .env 107 | .venv 108 | env/ 109 | venv/ 110 | ENV/ 111 | env.bak/ 112 | venv.bak/ 113 | 114 | # Spyder project settings 115 | .spyderproject 116 | .spyproject 117 | 118 | # Rope project settings 119 | .ropeproject 120 | 121 | # mkdocs documentation 122 | /site 123 | 124 | # mypy 125 | .mypy_cache/ 126 | .dmypy.json 127 | dmypy.json 128 | 129 | # Pyre type checker 130 | .pyre/ 131 | 132 | settings.json 133 | tutorial/config_example_6port_ms.py 134 | tutorial/config_example_6port.py 135 | 136 | test.py 137 | *.npy 138 | -------------------------------------------------------------------------------- /tutorial/convergence.py: -------------------------------------------------------------------------------- 1 | ''' Example tests showcasing the potential savings in computation time offered by AMR. The waveguide used in this test is a 3 port lantern. ''' 2 | 3 | import matplotlib.pyplot as plt 4 | import numpy as np 5 | from lightbeam import optics,LPmodes 6 | from lightbeam.mesh import RectMesh3D 7 | from lightbeam.misc import norm_nonu,normalize,overlap_nonu 8 | from lightbeam.prop import Prop3D 9 | 10 | 11 | from scipy.optimize import curve_fit 12 | outs = np.load('tutorial/convergence.npy') 13 | outs = np.vstack( (outs, np.array([0.30718035, 0.53222632, 0.14978008]))) 14 | times = np.array([45,67,99,236,665,2842,14447]) 15 | 16 | diffs = np.sqrt(np.sum(np.power(outs[:-1,:] - outs[-1,:],2),axis=1)) 17 | 18 | f = lambda x,a,b:a*x+b 19 | popt,pcov = curve_fit(f,np.log(times),np.log(diffs)) 20 | plt.figure(figsize=(4,3)) 21 | plt.plot(times,np.exp(f(np.log(times),*popt)),color='0.5',ls='dashed',lw=2) 22 | 23 | plt.plot(1115,0.000119196,marker='.',ms=9,color='steelblue',label='AMR',ls='None') 24 | plt.loglog(times,diffs,ls="None",marker='.',color='k',ms=9,label='uniform') 25 | plt.xlabel('computation time (s)') 26 | plt.ylabel('total error') 27 | plt.legend(frameon=False) 28 | plt.tight_layout() 29 | plt.show() 30 | 31 | def compute_port_power(ds,AMR=False,ref_val=2e-4,max_iters=5,remesh_every=50): 32 | ''' Compute the output port powers for a 3 port lantern, given some simulation parameters. ''' 33 | # wavelength 34 | wl = 1.0 #um 35 | 36 | # mesh 37 | xw = 64 #um 38 | yw= 64 #um 39 | zw = 10000 #um 40 | num_PML = int(4/ds) # number of cells 41 | dz = 1 42 | 43 | mesh = RectMesh3D(xw,yw,zw,ds,dz,num_PML) 44 | mesh.xy.max_iters = max_iters 45 | 46 | xg,yg = mesh.xg[num_PML:-num_PML,num_PML:-num_PML] , mesh.yg[num_PML:-num_PML,num_PML:-num_PML] 47 | 48 | # optic (3 port lantern) 49 | taper_factor = 4 50 | rcore = 2.2/taper_factor # INITIAL core radius 51 | rclad = 4 52 | nclad = 1.4504 53 | ncore = nclad + 0.0088 54 | njack = nclad - 5.5e-3 55 | 56 | lant = optics.lant3big(rcore,rclad,ncore,nclad,njack,rclad/2,zw,final_scale=taper_factor) 57 | 58 | def launch_field(x,y): 59 | return normalize(np.exp(10.j*x*wl/xw)*LPmodes.lpfield(x,y,0,1,rclad,wl,ncore,nclad)) 60 | 61 | # propagation 62 | 63 | prop = Prop3D(wl,mesh,lant,nclad) 64 | 65 | if AMR: 66 | u , u0 = prop.prop2end(launch_field,ref_val=ref_val,remesh_every=remesh_every) 67 | else: 68 | u = u0 = prop.prop2end_uniform(launch_field(xg,yg)) 69 | 70 | xg,yg = np.meshgrid(mesh.xy.xa,mesh.xy.ya,indexing='ij') 71 | 72 | w = mesh.xy.get_weights() 73 | 74 | # get the output port powers 75 | 76 | output_powers = [] 77 | for pos in lant.final_core_locs: 78 | _m = norm_nonu(LPmodes.lpfield(xg-pos[0],yg-pos[1],0,1,rcore*taper_factor,wl,ncore,nclad),w) 79 | output_powers.append(np.power(overlap_nonu(_m,u,w),2)) 80 | 81 | return np.array(output_powers) 82 | 83 | # example of how to use 84 | #output = compute_port_power(1/64,False) 85 | #print(output) 86 | -------------------------------------------------------------------------------- /tutorial/config_example.py: -------------------------------------------------------------------------------- 1 | ''' example configuration file for run_bpm_example.py ''' 2 | ################################ 3 | ## free space wavelength (um) ## 4 | ################################ 5 | 6 | wl0 = 1.55 7 | 8 | ######################## 9 | ## lantern parameters ## 10 | ######################## 11 | 12 | zex = 30000 # length of lantern, in um 13 | scale = 1/4 # how much smaller the input end is wrt the output end 14 | rcore = 4.5 * scale # how large the lantern cores are, at the input (um) 15 | rclad = 16.5 # how large the lantern cladding is, at the input (um) 16 | ncore = 1.4504 + 0.0088 # lantern core refractive index 17 | nclad = 1.4504 # cladding index 18 | njack = 1.4504 - 5.5e-3 # jacket index 19 | 20 | ################################### 21 | ## sampling grid parameters (um) ## 22 | ################################### 23 | 24 | xw0 = 128 # simulation zone x width (um) 25 | yw0 = 128 # simulation zone y width (um) 26 | zw = zex 27 | ds = 1 # base grid resolution (um) 28 | dz = 3 # z stepping resolution (um) 29 | 30 | ############################# 31 | ## mesh refinement options ## 32 | ############################# 33 | 34 | ref_val = 1e-4 # controls adaptive meshing. lower -> more careful 35 | remesh_every = 50 # how many z-steps to go before recomputing the adaptive mesh 36 | max_remesh_iters = 6 # maximum amount of subdivisions when computing the adaptive mesh 37 | 38 | xw_func = None # optional functions which allow the simulation zone to "grow" with z, which may save on computation time 39 | yw_func = None 40 | 41 | ################## 42 | ## PML settings ## 43 | ################## 44 | 45 | num_PML = 12 46 | sig_max = 3. + 0.j 47 | 48 | ###################### 49 | ## set launch field ## 50 | ###################### 51 | import numpy as np 52 | import matplotlib.pyplot as plt 53 | from lightbeam import LPmodes 54 | from lightbeam.misc import normalize 55 | 56 | xa = np.linspace(-xw0/2,xw0/2,int(xw0/ds)+1) 57 | ya = np.linspace(-yw0/2,yw0/2,int(yw0/ds)+1) 58 | xg,yg = np.meshgrid(xa,ya) 59 | 60 | u0 = normalize(LPmodes.lpfield(xg,yg,2,1,rclad,wl0,nclad,njack,'cos')) 61 | 62 | fplanewidth = 0 # manually reset the width of the input field. set to 0 to match field extent with grid extent. 63 | 64 | ##################### 65 | ## reference index ## 66 | ##################### 67 | 68 | n0 = 1.4504 69 | dynamic_n0 = False 70 | 71 | ################### 72 | ## monitor field ## 73 | ################### 74 | 75 | monitor_func = None 76 | 77 | ############################# 78 | ## write out field dist to ## 79 | ############################# 80 | 81 | writeto = None 82 | 83 | # generate optical element 84 | from lightbeam import optics 85 | optic = optics.lant19(rcore,rclad,ncore,nclad,njack,rclad/3,zex,final_scale=1/scale) 86 | 87 | 88 | ####################### 89 | ## initial core locs ## 90 | ####################### 91 | 92 | xpos_i = optic.core_locs[:,0] 93 | ypos_i = optic.core_locs[:,1] 94 | 95 | ##################### 96 | ## final core locs ## 97 | ##################### 98 | 99 | xpos = xpos_i / scale 100 | ypos = ypos_i / scale 101 | -------------------------------------------------------------------------------- /src/lightbeam/misc.py: -------------------------------------------------------------------------------- 1 | ''' bunch of miscellaneous functions that I didn't know where to put''' 2 | 3 | import numpy as np 4 | from bisect import bisect_left 5 | import time 6 | from scipy.interpolate import RectBivariateSpline 7 | 8 | def getslices(bounds,arr): 9 | '''given a range, get the idxs corresponding to that range in the sorted array arr ''' 10 | if len(bounds)==0: 11 | return np.s_[0:0] 12 | elif len(bounds)==1: 13 | return np.s_[bisect_left(arr,bounds[0])] 14 | elif len(bounds)==2: 15 | return np.s_[bisect_left(arr,bounds[0]):bisect_left(arr,bounds[1])+1] 16 | else: 17 | raise Exception("malformed bounds input in getslices(); check savex,savey,savez in config.py") 18 | 19 | def resize2(image,newshape): 20 | '''another resampling function that uses scipy, not cv2''' 21 | xpix = np.arange(image.shape[0]) 22 | ypix = np.arange(image.shape[1]) 23 | 24 | xpix_new = np.linspace(xpix[0],xpix[-1],newshape[0]) 25 | ypix_new = np.linspace(ypix[0],ypix[-1],newshape[1]) 26 | 27 | return RectBivariateSpline(xpix,ypix,image)(xpix_new,ypix_new) 28 | 29 | def overlap(u1,u2,weight=1,c=False): 30 | if not c: 31 | return weight*np.abs(np.sum(np.conj(u1)*u2)) 32 | return weight * np.sum(np.conj(u1)*u2) 33 | 34 | def overlap_nonu(u1,u2,weights): 35 | return np.abs(np.sum(weights*np.conj(u1)*u2)) 36 | 37 | def overlap_nonu_trap(u1,u2,xa,ya,c=False): 38 | integrand = np.conj(u1)*u2 39 | integral = np.traps(np.trapz(integrand,ya,axis=-1),xa) 40 | if c: 41 | return integral 42 | return np.abs(integral) 43 | 44 | def normalize(u0,weight=1,normval = 1): 45 | norm = np.sqrt(normval/overlap(u0,u0,weight)) 46 | u0 *= norm 47 | return u0 48 | 49 | def norm_nonu(u0,weights,normval = 1): 50 | norm = np.sqrt(normval/overlap_nonu(u0,u0,weights)) 51 | u0 *= norm 52 | return u0 53 | 54 | def printProgressBar (iteration, total, prefix = '', suffix = '', decimals = 1, length = 100, fill = '█', printEnd = "\r"): 55 | """ 56 | Pulled from https://stackoverflow.com/questions/3173320/text-progress-bar-in-the-console 57 | 58 | Call in a loop to create terminal progress bar 59 | @params: 60 | iteration - Required : current iteration (Int) 61 | total - Required : total iterations (Int) 62 | prefix - Optional : prefix string (Str) 63 | suffix - Optional : suffix string (Str) 64 | decimals - Optional : positive number of decimals in percent complete (Int) 65 | length - Optional : character length of bar (Int) 66 | fill - Optional : bar fill character (Str) 67 | printEnd - Optional : end character (e.g. "\r", "\r\n") (Str) 68 | """ 69 | percent = ("{0:." + str(decimals) + "f}").format(100 * (iteration / float(total))) 70 | filledLength = int(length * iteration // total) 71 | bar = fill * filledLength + '-' * (length - filledLength) 72 | print('\r%s |%s| %s%% %s' % (prefix, bar, percent, suffix), end = printEnd) 73 | # Print New Line on Complete 74 | if iteration == total: 75 | print() 76 | 77 | def timeit(method): 78 | '''pulled from someone's github or something. can't find it anymore''' 79 | def timed(*args, **kw): 80 | ts = time.time() 81 | result = method(*args, **kw) 82 | te = time.time() 83 | if 'log_time' in kw: 84 | name = kw.get('log_name', method.__name__.upper()) 85 | kw['log_time'][name] = int((te - ts)) 86 | else: 87 | print('%r %2.4f s' % \ 88 | (method.__name__, (te - ts))) 89 | return result 90 | return timed 91 | 92 | def gauss(xg,yg,theta,phi,sigu,sigv,k0,x0=0,y0=0.): 93 | '''tilted gaussian beam''' 94 | u = np.cos(theta)*np.cos(phi)*(xg-x0) + np.cos(theta)*np.sin(phi)*(yg-y0) 95 | v = -np.sin(phi)*(xg-x0) + np.cos(phi)*(yg-y0) 96 | w = np.sin(theta)*np.cos(phi)*(xg-x0) + np.sin(theta)*np.sin(phi)*(yg-y0) 97 | out = ( np.exp(1.j*k0*w)*np.exp(-0.5*np.power(u/sigu,2.)-0.5*np.power(v/sigv,2.)) ).astype(np.complex128) 98 | return out/np.sqrt(overlap(out,out)) 99 | 100 | def read_rsoft(fname): 101 | arr = np.loadtxt(fname,skiprows = 4).T 102 | reals = arr[::2] 103 | imags = arr[1::2] 104 | field = (reals+1.j*imags).T 105 | return field.astype(np.complex128) 106 | 107 | def write_rsoft(fname,u0,xw,yw): 108 | '''save field to a file format useable by rsoft''' 109 | out = np.empty((u0.shape[0]*2,u0.shape[1])) 110 | reals = np.real(u0) 111 | imags = np.imag(u0) 112 | 113 | for j in range(out.shape[0]): 114 | if j%2==0: 115 | out[j] = reals[:,int(j/2)] 116 | else: 117 | out[j] = imags[:,int(j/2)] 118 | 119 | header = "/rn,a,b/nx0/ls1\n/r,qa,qb\n{} {} {} 0 OUTPUT_REAL_IMAG_3D\n{} {} {}".format(u0.shape[0],-xw/2,xw/2,u0.shape[1],-yw/2,yw/2) 120 | np.savetxt(fname+".dat", out.T, header = header, fmt = "%f",comments="",newline="\n") -------------------------------------------------------------------------------- /src/lightbeam/zernike.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from scipy.special import factorial,jn,gamma 3 | import matplotlib.pyplot as plt 4 | from numpy import fft 5 | from numpy.polynomial import polynomial as poly 6 | 7 | """ functions for working with Zernike modes """ 8 | 9 | def nm2j(n,m): 10 | """ convert n,m indices to j indices following Noll's prescription """ 11 | offset = 0 12 | if ( m >= 0 and (n%4 in [2,3]) ) or ( m<=0 and (n%4 in [0,1]) ): 13 | offset = 1 14 | 15 | j = int( n*(n+1)/2 ) + abs(m) + offset 16 | return j 17 | 18 | def j2nm(j): 19 | """ inverse of nm2j """ 20 | 21 | _nmin = int (0.5 * (-3 + np.sqrt(8*j+1))) 22 | _nmax = int (0.5 * (-1 + np.sqrt(8*j+1))) 23 | 24 | _m1 = j - int(_nmin * (_nmin+1)/2) - 1 25 | _m2 = _m1 + 1 26 | _m3 = j - int(_nmax * (_nmax+1)/2) - 1 27 | _m4 = _m3 + 1 28 | 29 | guesses = [ [_nmin,_m1], [_nmin,_m2], [_nmax,_m3], [_nmax,_m4] ] 30 | for n,m in guesses: 31 | if m<0 or (n-m)%2 != 0 or m>n: 32 | continue 33 | if j%2 == 0: 34 | return (n,m) 35 | else: 36 | return (n,-m) 37 | 38 | raise Exception("CATASTROPIC FAILURE IN j2nm()") 39 | 40 | def Z_rad(n,m): 41 | """radial component of Zernike mode""" 42 | m = abs(m) 43 | indices = np.arange(0,1+int((n-m)/2)) 44 | coeffs = np.power(-1,indices) * factorial(n-indices) / factorial(indices) / factorial( int((n+m)/2) - indices ) / factorial( int((n-m)/2) - indices) 45 | powers = n - 2 * indices 46 | 47 | def _inner_(r): 48 | return np.sum( coeffs[:,None,None] * np.power( np.repeat([r],len(powers),axis=0) , powers[:,None,None] ) , axis = 0 ) 49 | 50 | return _inner_ 51 | 52 | def Z_az(n,m): 53 | """azimuthal component of Zernike mode""" 54 | def _inner_(theta): 55 | if m == 0: 56 | return np.sqrt(n+1) 57 | elif m>0: 58 | return np.sqrt(2*(n+1)) * np.cos( m*theta ) 59 | else: 60 | return np.sqrt(2*(n+1)) * np.sin( abs(m)*theta ) 61 | return _inner_ 62 | 63 | def Z(n,m): 64 | """Zernike mode (n,m indexed). returns function of r,theta""" 65 | assert (n-abs(m))%2 == 0 , "in Z(n,m), n - |m| must be even" 66 | def _inner_(r,theta): 67 | return ( Z_rad(n,m)(r) * Z_az(n,m)(theta) ) * (r<=1) 68 | return _inner_ 69 | 70 | def Zj_pol(j): 71 | """Zernike mode (j indexed). returns function of r,theta""" 72 | n,m = j2nm(j) 73 | return Z(n,m) 74 | 75 | def Zj_cart(j): 76 | """Zernike mode (j indexed). returns function of x,y""" 77 | def _inner_(x,y): 78 | r = np.sqrt(x*x+y*y) 79 | t = np.arctan2(y,x) 80 | return Zj_pol(j)(r,t) 81 | return _inner_ 82 | 83 | def Qj(j): 84 | """Fourier transform of Zernike mode""" 85 | n,m = j2nm(j) 86 | if j == 1: 87 | lim = 1 88 | else: 89 | lim = 0 90 | 91 | def _inner_(k,phi): 92 | with np.errstate(divide="ignore",invalid="ignore"): 93 | kdep = np.where( k!=0 , jn(n+1, 2*np.pi*k) / (np.pi*k),lim) 94 | if m>0: 95 | return np.sqrt(2*(n+1)) * kdep * np.power(-1,int((n-m)/2)) * np.power(-1j,m) * np.cos(m*phi) 96 | elif m<0: 97 | return np.sqrt(2*(n+1)) * kdep * np.power(-1,int((n-m)/2)) * np.power(-1j,m) * np.sin(m*phi) 98 | else: 99 | return np.sqrt(n+1) * kdep * np.power(-1,int(n/2)) 100 | return _inner_ 101 | 102 | def Qj_cart(j): 103 | def _inner_(kx,ky): 104 | k = np.sqrt(kx*kx+ky*ky) 105 | phi = np.arctan2(ky,kx) 106 | return Qj(j)(k,phi) 107 | return _inner_ 108 | 109 | def high_pass(n): 110 | """ returns a high pass filter that can be applied to any turbulence PSD 111 | to remove Zernike modes up to and including order n 112 | """ 113 | def _inner_(kx,ky): 114 | filt = np.zeros_like(kx) 115 | for j in range(1,n+1): 116 | filt += np.power(np.abs(Qj_cart(j)(kx,ky)),2) 117 | return 1 - filt 118 | return _inner_ 119 | 120 | def low_pass(n): 121 | """ returns a low pass filter that can be applied to any turbulence PSD 122 | to only include Zernike modes up to and including order n """ 123 | def _inner_(kx,ky): 124 | filt = np.zeros_like(kx) 125 | for j in range(1,n+1): 126 | filt += np.power(np.abs(Qj_cart(j)(kx,ky)),2) 127 | return filt 128 | return _inner_ 129 | 130 | def noll_cov(i,j,D,r0): 131 | """covariance matrix for Kolmogorov turbulence""" 132 | 133 | n,m = j2nm(i) 134 | _n,_m = j2nm(j) 135 | 136 | if m!= _m or (i-j)%2==1: 137 | return 0 138 | 139 | Kzz_fac = gamma(14/3)* np.power(24/5 * gamma(6/5),5/6) * np.power(gamma(11/6),2) / (2*np.pi**2) 140 | Kzz = Kzz_fac * np.power(-1, int((n+_n-2*abs(m))/2) ) * np.sqrt((n+1)*(_n+1)) 141 | fac1 = 0.5 * (n + _n - 5/3) 142 | fac2 = 0.5 * (n - _n + 17/3) 143 | fac3 = 0.5 * (_n - n + 17/3) 144 | fac4 = 0.5 * (n + _n + 23/3) 145 | 146 | return ( Kzz * gamma(fac1) * np.power(D/r0,5/3) ) / (gamma(fac2)*gamma(fac3)*gamma(fac4)) 147 | 148 | def compute_noll_mat(N,save=True): 149 | """ compute normalized noll matrix in (D/r0)^5/3 units 150 | args 151 | 152 | N: number of Zernike modes to use 153 | save: flag to save matrix to "nollmat.npy" 154 | 155 | returns: Zernike covariance matrix 156 | """ 157 | mat = np.empty((N,N)) 158 | for i in range(2,N+2): 159 | for j in range(2,N+2): 160 | mat[i-2,j-2] = noll_cov(i,j,1,1) 161 | if save: 162 | np.save("nollmat",mat) 163 | return mat 164 | 165 | def phase_screen_func(tmat): 166 | N = tmat.shape[0] 167 | a = np.random.normal(size=N) 168 | a = np.dot(tmat,a) 169 | 170 | def _inner_(x,y): 171 | out = np.zeros_like(x) 172 | for i in range(N): 173 | j = i+2 174 | out += a[i]*Zj_cart(j)(x,y) 175 | return out 176 | 177 | return _inner_ 178 | 179 | def inner_product(u0,u1,ds): 180 | """inner product which makes use of orthogonality relation of Zernikes as defined in this file""" 181 | return np.sum(u0*u1) * ds * ds / np.pi 182 | 183 | ### some aliases for more consistent naming ### 184 | 185 | def zkfield(xg,yg,j): 186 | """compute a Zernike mode over a rectangular grid 187 | ARGS: 188 | xg: x coordinate grid (2D) 189 | yg: y coordinate grid (2D) 190 | j: noll index of Zernike mode 191 | """ 192 | r = np.sqrt(xg*xg+yg*yg) 193 | t = np.arctan2(yg,xg) 194 | return Zj_pol(j)(r,t) -------------------------------------------------------------------------------- /src/lightbeam/geom.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from numpy import logical_not as NOT, logical_and as AND, logical_or as OR 3 | from numba import njit 4 | 5 | '''a collection of functions for antialiasing circles''' 6 | 7 | ## calculate circle-square overlap. 8 | 9 | # original code in pixwt.c by Marc Buie 10 | # ported to pixwt.pro (IDL) by Doug Loucks, Lowell Observatory, 1992 Sep 11 | # subsequently ported to python by Michael Fitzgerald, 12 | # LLNL, fitzgerald15@llnl.gov, 2007-10-16 13 | 14 | ### Marc Buie, you are my hero 15 | 16 | def _arc(x, y0, y1, r): 17 | """ 18 | Compute the area within an arc of a circle. The arc is defined by 19 | the two points (x,y0) and (x,y1) in the following manner: The 20 | circle is of radius r and is positioned at the originp. The origin 21 | and each individual point define a line which intersects the 22 | circle at some point. The angle between these two points on the 23 | circle measured from y0 to y1 defines the sides of a wedge of the 24 | circle. The area returned is the area of this wedge. If the area 25 | is traversed clockwise then the area is negative, otherwise it is 26 | positive. 27 | """ 28 | thetas = np.empty_like(x) 29 | mask = (x==0) 30 | thetas[mask] = np.pi/2*(np.sign(y1[mask])-np.sign(y0[mask])) 31 | thetas[~mask] = np.arctan(y1[~mask]/x[~mask])-np.arctan(y0[~mask]/x[~mask]) 32 | return 0.5*r**2*thetas 33 | #thetas = np.where(x==0, np.pi/2*(np.sign(y1)-np.sign(y0)),np.arctan(y1/x)-np.arctan(y0/x)) 34 | #return 0.5 * r**2 * (np.arctan2(y1,x) - np.arctan2(y0,x)) 35 | 36 | def _chord(x, y0, y1): 37 | """ 38 | Compute the area of a triangle defined by the origin and two 39 | points, (x,y0) and (x,y1). This is a signed area. If y1 > y0 40 | then the area will be positive, otherwise it will be negative. 41 | """ 42 | return 0.5 * x * (y1 - y0) 43 | 44 | def _oneside(x, y0, y1, r): 45 | """ 46 | Compute the area of intersection between a triangle and a circle. 47 | The circle is centered at the origin and has a radius of r. The 48 | triangle has verticies at the origin and at (x,y0) and (x,y1). 49 | This is a signed area. The path is traversed from y0 to y1. If 50 | this path takes you clockwise the area will be negative. 51 | """ 52 | 53 | if np.all((x==0)): return x 54 | 55 | sx = x.shape 56 | ans = np.zeros(sx, dtype=np.float64) 57 | yh = np.zeros(sx, dtype=np.float64) 58 | to = (abs(x) >= r) 59 | ti = (abs(x) < r) 60 | if np.any(to): 61 | ans[to] = _arc(x[to], y0[to], y1[to], r) 62 | if not np.any(ti): 63 | return ans 64 | 65 | yh[ti] = np.sqrt(r**2 - x[ti]**2) 66 | 67 | i = ((y0 <= -yh) & ti) 68 | if np.any(i): 69 | 70 | j = ((y1 <= -yh) & i) 71 | if np.any(j): 72 | ans[j] = _arc(x[j], y0[j], y1[j], r) 73 | 74 | j = ((y1 > -yh) & (y1 <= yh) & i) 75 | if np.any(j): 76 | ans[j] = _arc(x[j], y0[j], -yh[j], r) + \ 77 | _chord(x[j], -yh[j], y1[j]) 78 | 79 | j = ((y1 > yh) & i) 80 | if np.any(j): 81 | ans[j] = _arc(x[j], y0[j], -yh[j], r) + \ 82 | _chord(x[j], -yh[j], yh[j]) + \ 83 | _arc(x[j], yh[j], y1[j], r) 84 | 85 | i = ((y0 > -yh) & (y0 < yh) & ti) 86 | if np.any(i): 87 | 88 | j = ((y1 <= -yh) & i) 89 | if np.any(j): 90 | ans[j] = _chord(x[j], y0[j], -yh[j]) + \ 91 | _arc(x[j], -yh[j], y1[j], r) 92 | 93 | j = ((y1 > -yh) & (y1 <= yh) & i) 94 | if np.any(j): 95 | ans[j] = _chord(x[j], y0[j], y1[j]) 96 | 97 | j = ((y1 > yh) & i) 98 | if np.any(j): 99 | ans[j] = _chord(x[j], y0[j], yh[j]) + \ 100 | _arc(x[j], yh[j], y1[j], r) 101 | 102 | i = ((y0 >= yh) & ti) 103 | if np.any(i): 104 | 105 | j = ((y1 <= -yh) & i) 106 | if np.any(j): 107 | ans[j] = _arc(x[j], y0[j], yh[j], r) + \ 108 | _chord(x[j], yh[j], -yh[j]) + \ 109 | _arc(x[j], -yh[j], y1[j], r) 110 | 111 | j = ((y1 > -yh) & (y1 <= yh) & i) 112 | if np.any(j): 113 | ans[j] = _arc(x[j], y0[j], yh[j], r) + \ 114 | _chord(x[j], yh[j], y1[j]) 115 | 116 | j = ((y1 > yh) & i) 117 | if np.any(j): 118 | ans[j] = _arc(x[j], y0[j], y1[j], r) 119 | 120 | return ans 121 | 122 | def _intarea(xc, yc, r, x0, x1, y0, y1): 123 | """ 124 | Compute the area of overlap of a circle and a rectangle. 125 | xc, yc : Center of the circle. 126 | r : Radius of the circle. 127 | x0, y0 : Corner of the rectangle. 128 | x1, y1 : Opposite corner of the rectangle. 129 | """ 130 | x0 = x0 - xc 131 | y0 = y0 - yc 132 | x1 = x1 - xc 133 | y1 = y1 - yc 134 | return _oneside(x1, y0, y1, r) + _oneside(y1, -x1, -x0, r) + \ 135 | _oneside(-x0, -y1, -y0, r) + _oneside(-y0, x0, x1, r) 136 | 137 | def nonu_pixwt(xc,yc,r,x,y,rx,ry,dx,dy): 138 | 139 | area = dx*dy/4*(rx+1)*(ry+1) 140 | return _intarea(xc,yc,r,x-0.5*dx,x+0.5*rx*dx,y-0.5*dy,y+0.5*ry*dy)/area 141 | 142 | def AA_circle_nonu(out,xg,yg,xgh,ygh,center,R,n0,n1,where,rxg,ryg,dxg,dyg): 143 | 144 | xdif = xgh-center[0] 145 | ydif = ygh-center[1] 146 | rsqh = xdif*xdif+ydif*ydif 147 | mask_in,mask_b = get_masks(rsqh,R*R) 148 | out[where][mask_in] = n1 149 | x0,y0 = xg[mask_b],yg[mask_b] 150 | 151 | rx,ry = rxg[where][mask_b], ryg[where][mask_b] 152 | dx,dy = dxg[where][mask_b], dyg[where][mask_b] 153 | 154 | area = nonu_pixwt(center[0],center[1],R,x0,y0,rx,ry,dx,dy) 155 | newvals = n1*area+n0*(1-area) 156 | out[where][mask_b] = newvals 157 | 158 | def pixwt(xc, yc, r, x, y): 159 | """ 160 | Compute the fraction of a unit pixel that is interior to a circle. 161 | The circle has a radius r and is centered at (xc, yc). The center 162 | of the unit pixel (length of sides = 1) is at (x, y). 163 | 164 | Divides the circle and rectangle into a series of sectors and 165 | triangles. Determines which of nine possible cases for the 166 | overlap applies and sums the areas of the corresponding sectors 167 | and triangles. 168 | 169 | area = pixwt( xc, yc, r, x, y ) 170 | 171 | xc, yc : Center of the circle, numeric scalars 172 | r : Radius of the circle, numeric scalars 173 | x, y : Center of the unit pixel, numeric scalar or vector 174 | """ 175 | return _intarea(xc, yc, r, x-0.5, x+0.5, y-0.5, y+0.5) 176 | 177 | def get_masks(rsqh,R2): 178 | maskh = np.full(rsqh.shape,False) 179 | maskh[rsqh<=R2] = True 180 | 181 | mask_in = AND(maskh[1:,1:], AND(maskh[1:,:-1], AND(maskh[:-1,1:],maskh[:-1,:-1]))) 182 | mask_out = AND(NOT(maskh[1:,1:]), AND(NOT(maskh[1:,:-1]), AND(NOT(maskh[:-1,1:]),NOT(maskh[:-1,:-1])))) 183 | mask_b = NOT(OR(mask_in,mask_out)) 184 | 185 | return mask_in,mask_b 186 | 187 | def AA_circle(out,xg,yg,xgh,ygh,center,R,n0,n1,ds,where): 188 | xdif = xgh-center[0] 189 | ydif = ygh-center[1] 190 | rsqh = xdif*xdif+ydif*ydif 191 | mask_in,mask_b = get_masks(rsqh,R*R) 192 | out[where][mask_in] = n1 193 | x0,y0 = xg[mask_b],yg[mask_b] 194 | area = pixwt(center[0]/ds,center[1]/ds,R/ds,x0/ds,y0/ds) 195 | newvals = n1*area+n0*(1-area) 196 | out[where][mask_b] = newvals 197 | 198 | 199 | 200 | 201 | 202 | 203 | 204 | 205 | 206 | 207 | 208 | 209 | 210 | -------------------------------------------------------------------------------- /src/lightbeam/screen.py: -------------------------------------------------------------------------------- 1 | # Michael Fitzgerald (mpfitz@ucla.edu) 2 | import numpy as np 3 | from numpy.fft import fftfreq, fft2, ifft2,fftshift 4 | import matplotlib.pyplot as plt 5 | from matplotlib import animation 6 | 7 | def boiling_freq(k,t1): 8 | return np.power(k,2/3)/t1 9 | 10 | # from Srinath et al. (2015) 28 Dec 2015 | Vol. 23, No. 26 | DOI:10.1364/OE.23.033335 | OPTICS EXPRESS 33335 11 | class PhaseScreenGenerator(object): 12 | def __init__(self, D, p, vy, vx, T, r0, wl0, wl, rs=None, seed=None, alpha_mag=1.,filter_func = None,filter_scale=None): 13 | # set up random number generator 14 | if rs is None: 15 | rs = np.random.RandomState(seed=seed) 16 | self.seed = seed 17 | 18 | if filter_scale is None: 19 | filter_scale = D/2 20 | 21 | if filter_func is None: 22 | filter_func = lambda x,y: np.ones_like(x) 23 | 24 | self.rs = rs 25 | 26 | # set up array dimensions 27 | a = 2.**np.ceil(np.log2(D/p)) / (D/p) 28 | self.N = int(a*D/p) # number of pixels on a side of a screen 29 | S = self.N*p # [m] screen size 30 | 31 | # frequency array 32 | fy = fftfreq(self.N, d=p) 33 | fx = fftfreq(self.N, d=p) 34 | 35 | fy = np.dot(fy[:,np.newaxis], np.ones(self.N)[np.newaxis,:]) 36 | fx = np.dot(np.ones(self.N)[:,np.newaxis], fx[np.newaxis,:]) 37 | 38 | ff = np.sqrt(fx*fx+fy*fy) 39 | 40 | # turbulent power spectrum 41 | with np.errstate(divide='ignore'): 42 | self.P = 2.*np.pi/S * self.N * r0**(-5./6.) * (fy*fy + fx*fx)**(-11./12.) * np.sqrt(0.00058) * (wl0/wl) 43 | self.P[0,0] = 0. # no infinite power 44 | 45 | if filter_func is not None: 46 | self.P *= np.sqrt(filter_func(filter_scale*fx,filter_scale*fy)) 47 | """ 48 | plt.loglog(fftshift(fx[0]), np.sqrt(fftshift(filter_func(filter_scale*fx[0],0)))*np.max(self.P[0]) ,color='white') 49 | 50 | for _fx,_P in zip(fx[0],self.P[0]): 51 | plt.plot((_fx,_fx),(0,_P),color='0.5',lw=1) 52 | plt.loglog(fx[0],self.P[0],marker='.',markerfacecolor='steelblue',markeredgecolor='white',markersize=10,ls='None',markeredgewidth=0.5) 53 | 54 | 55 | 56 | plt.loglog(fx[0],self.P[0],marker='.',markerfacecolor='indianred',markeredgecolor='white',markersize=10,ls='None',markeredgewidth=0.5) 57 | plt.xlabel(r"wavenumber ($m^{-1}$) ") 58 | plt.ylabel("power") 59 | plt.show() 60 | """ 61 | 62 | 63 | # set phase scale 64 | theta = -2.*np.pi*T*(fy*vy+fx*vx) 65 | self.alpha = alpha_mag*np.exp(1j*theta)#*np.exp(-T*boiling_freq(ff,0.27)) # |alpha|=1 is pure frozen flow 66 | 67 | # white noise scale factor 68 | self.beta = np.sqrt(1.-np.abs(self.alpha)**2)*self.P 69 | 70 | self.last_phip = None 71 | 72 | self.t = 0 73 | 74 | def generate(self): 75 | # generate white noise 76 | 77 | w = self.rs.randn(self.N,self.N) 78 | wp = fft2(w) 79 | 80 | # get FT of phase 81 | if self.last_phip is None: 82 | # first timestep, generate power 83 | phip = self.P*wp 84 | else: 85 | phip = self.alpha*self.last_phip + self.beta*wp 86 | self.last_phip = phip 87 | 88 | # get phase 89 | phi = ifft2(phip).real 90 | 91 | return phi 92 | 93 | def reset(self): 94 | self.last_phip = None 95 | self.rs = np.random.RandomState(self.seed) 96 | 97 | def get_layer(self,t): 98 | # need to generate from self.t up until t in steps of T 99 | pass 100 | 101 | def make_ani(out,t,dt,D=10,p=0.1,wl0=1,wl=1,vx=4,vy=0,r0=0.25,alpha=1,seed=345698,delay=20): 102 | psgen = PhaseScreenGenerator(D, p, vy, vx, dt, r0, wl0, wl, seed=seed,alpha_mag=alpha) 103 | 104 | # First set up the figure, the axis, and the plot element we want to animate 105 | fig = plt.figure() 106 | ax = plt.axes(xlim=(0, 128), ylim=(0,128)) 107 | 108 | im=plt.imshow(psgen.generate()) 109 | 110 | # initialization function: plot the background of each frame 111 | def init(): 112 | im.set_data(np.zeros((128,128))) 113 | return [im] 114 | 115 | # animation function. This is called sequentially 116 | def animate(i): 117 | im.set_data(psgen.generate()) 118 | return [im] 119 | 120 | # call the animator. blit=True means only re-draw the parts that have changed. 121 | anim = animation.FuncAnimation(fig, animate, init_func=init, 122 | frames=int(t/dt), interval=20, blit=True) 123 | 124 | anim.save(out+'.mp4', fps=60, extra_args=['-vcodec', 'libx264']) 125 | 126 | plt.show() 127 | 128 | if __name__=='__main__': 129 | 130 | import zernike as zk 131 | #make_ani("turb_ani_test_alpha0pt999",5,0.01,vx=0,alpha=0.99) 132 | 133 | r0s = [0.169,0.1765,0.185,0.196,0.215,0.245,0.295] 134 | SRs = [0.1,0.2,0.3,0.4,0.5,0.6,0.7] 135 | 136 | rms2s=[] 137 | rms3s=[] 138 | 139 | for sr,r0 in zip(SRs,r0s): 140 | 141 | D = 10. # [m] telescope diameter 142 | p = 10/256 # [m/pix] sampling scale 143 | 144 | # set wind parameters 145 | vy, vx = 0., 10. # [m/s] wind velocity vector 146 | T = 0.01 # [s] sampling interval 147 | 148 | # set turbulence parameters 149 | #r0 = 0.185 # [m] 150 | wl0 = 1 #[um] 151 | wl = 1 #[um] 152 | 153 | seed = 123456 154 | psgen = PhaseScreenGenerator(D, p, vy, vx, T, r0, wl0, wl, seed=seed) 155 | 156 | xa = ya = np.linspace(-1,1,256) 157 | xg , yg = np.meshgrid(xa,ya) 158 | 159 | dA = (D/256)**2 160 | 161 | z2 = zk.Zj_cart(2)(xg,yg) 162 | z3 = zk.Zj_cart(3)(xg,yg) 163 | 164 | #plt.imshow(z2) 165 | #plt.show() 166 | 167 | print(np.sum(z2*z2) * dA / np.pi/25) 168 | s = psgen.generate() 169 | 170 | 171 | PV2s = [] 172 | PV3s = [] 173 | c2s = [] 174 | c3s = [] 175 | pe2s = [] 176 | pe3s = [] 177 | 178 | for i in range(1000): 179 | s = psgen.generate() 180 | #plt.imshow(s) 181 | #plt.show() 182 | 183 | c2 = np.sum(s*z2)*dA/np.pi/25 184 | c3 = np.sum(s*z3)*dA/np.pi/25 185 | 186 | c2s.append(c2) 187 | c3s.append(c3) 188 | 189 | PV2s.append(4*c2/(2*np.pi)) 190 | PV3s.append(4*c3/(2*np.pi)) 191 | 192 | pe2 = np.std(c2*z2)/(2*np.pi) 193 | pe3 = np.std(c3*z3)/(2*np.pi) 194 | 195 | pe2s.append(pe2) 196 | pe3s.append(pe3) 197 | 198 | print(np.std(np.array(c2s))) 199 | plt.plot(np.arange(1000)*0.01,c2s) 200 | plt.show() 201 | 202 | rms2 = np.sqrt(np.mean(np.power(np.array(PV2s),2))) 203 | rms3 = np.sqrt(np.mean(np.power(np.array(PV3s),2))) 204 | 205 | plt.plot(np.arange(100),pe2s,color='steelblue',label='x tilt') 206 | plt.plot(np.arange(100),pe3s,color='indianred',label='y tilt') 207 | #plt.axhline(y=rms2,color='steelblue',ls='dashed') 208 | #plt.axhline(y=rms3,color='indianred',ls='dashed') 209 | plt.xlabel("timestep") 210 | plt.ylabel("RMS wf error of TT modes") 211 | 212 | plt.legend(frameon=False) 213 | plt.show() 214 | 215 | rms2s.append(rms2) 216 | rms3s.append(rms3) 217 | 218 | plt.plot(SRs,rms2s,label="x tilt",color='steelblue') 219 | plt.plot(SRs,rms3s,label="y tilt",color='indianred') 220 | 221 | plt.xlabel("Strehl ratio") 222 | plt.ylabel("rms mode amplitude") 223 | plt.legend(frameon=False) 224 | plt.show() 225 | 226 | """ 227 | plt.plot(np.arange(len(coeffs2)),coeffs2,color='steelblue',label="x tilt") 228 | plt.plot(np.arange(len(coeffs3)),coeffs3,color='indianred',label="y tilt") 229 | 230 | rms2 = np.sqrt(np.mean(np.power(np.array(coeffs2),2))) 231 | rms3 = np.sqrt(np.mean(np.power(np.array(coeffs3),2))) 232 | 233 | plt.axhline(y=rms2,color='steelblue',ls='dashed') 234 | plt.axhline(y=rms3,color='indianred',ls='dashed') 235 | plt.xlabel("timestep") 236 | plt.ylabel("zernike mode amplitude") 237 | plt.ylim(-8,8) 238 | 239 | plt.legend(frameon=False) 240 | 241 | plt.show() 242 | """ 243 | 244 | """ 245 | last = psgen.generate() 246 | out = np.zeros_like(last) 247 | 248 | for i in range(8000): 249 | cur = psgen.generate() 250 | out += np.power(cur-last,2) 251 | last = cur 252 | 253 | out/=8000 254 | plt.imshow(out) 255 | plt.show() 256 | print(np.mean(out)) 257 | """ 258 | 259 | 260 | 261 | #fits.writeto('screens_test.fits', screens, overwrite=True) 262 | 263 | 264 | 265 | 266 | ## show results 267 | #import pylab 268 | #fig = pylab.figure(0) 269 | #fig.clear() 270 | #ax = fig.add_subplot(111) 271 | #for screen in screens: 272 | # ax.cla() 273 | # ax.imshow(screen) 274 | # pylab.draw() 275 | # pylab.show() 276 | -------------------------------------------------------------------------------- /src/lightbeam/LPmodes.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from scipy.special import jn_zeros, jn, kn,jv,kv 3 | from scipy.optimize import brentq 4 | from numpy.lib import scimath 5 | 6 | def get_NA(ncore,nclad): 7 | return np.sqrt(ncore*ncore - nclad*nclad) 8 | 9 | def get_V(k0,rcore,ncore,nclad): 10 | return k0 * rcore * get_NA(ncore,nclad) 11 | 12 | def get_MFD(k0,rcore,ncore,nclad): 13 | """Marcuse approx. for straight step index fiber""" 14 | V = get_V(k0,rcore,ncore,nclad) 15 | return 2 * rcore * (0.65 + 1.619/np.power(V,1.5) + 2.879/np.power(V,6)) 16 | 17 | def get_MFD_from_NA(k0,rcore,ncore,NA): 18 | nclad = np.sqrt(ncore*ncore - NA*NA) 19 | return get_MFD(k0,rcore,ncore,nclad) 20 | 21 | def get_modes(V): 22 | '''frequency cutoff occurs when b(V) = 0. solve eqn 4.19. 23 | checks out w/ the function in the fiber ipynb''' 24 | 25 | l = 0 26 | m = 1 27 | modes = [] 28 | while True: 29 | 30 | if l == 0: 31 | #solving dispersion relation leads us to the zeros of J_1 32 | #1st root of J_1 is 0. 33 | 34 | modes.append((0,1)) 35 | 36 | while jn_zeros(1,m)[m-1]< V: 37 | 38 | modes.append((l,m+1)) 39 | m+=1 40 | else: 41 | #solving dispersion relation leads us to the zeros of J_l-1, neglecting 0 42 | if jn_zeros(l-1,1)[0]>V: 43 | break 44 | 45 | while jn_zeros(l-1,m)[m-1] 0: 55 | return jn_zeros(l-1, mmax) 56 | else: 57 | if mmax > 1: 58 | return np.concatenate(((0.,),jn_zeros(l-1, mmax-1))) 59 | else: 60 | return np.array((0.,)) 61 | 62 | def findBetween(solve_fn, lowbound, highbound, args=(), maxj=10): 63 | v = [lowbound, highbound] 64 | s = [solve_fn(lowbound, *args), solve_fn(highbound, *args)] 65 | 66 | if s[0] == 0.: return lowbound 67 | if s[1] == 0.: return highbound 68 | 69 | from itertools import count 70 | for j in count(): # probably not needed... 71 | if j == maxj: 72 | print("findBetween: max iter reached") 73 | return v[np.argmin(np.abs(s))] 74 | #return np.nan 75 | for i in range(len(s)-1): 76 | a, b = v[i], v[i+1] 77 | fa, fb = s[i], s[i+1] 78 | 79 | if (fa > 0 and fb < 0) or (fa < 0 and fb > 0): 80 | z = brentq(solve_fn, a, b, args=args) 81 | fz = solve_fn(z, *args) 82 | if abs(fa) > abs(fz) < abs(fb): # Skip discontinuities 83 | return z 84 | 85 | ls = len(s) 86 | for i in range(ls-1): 87 | a, b = v[2*i], v[2*i+1] 88 | c = (a + b) / 2 89 | v.insert(2*i+1, c) 90 | s.insert(2*i+1, solve_fn(c, *args)) 91 | 92 | def get_b(l, m, V): 93 | if l == 0: 94 | def solve_fn(b, V): 95 | v = V*np.sqrt(b) 96 | u = V*np.sqrt(1.-b) 97 | return (u * jn(1, u) * kn(0, v) - v * jn(0, u) * kn(1, v)) 98 | else: 99 | def solve_fn(b, V): 100 | v = V*np.sqrt(b) 101 | u = V*np.sqrt(1.-b) 102 | return (u * jn(l - 1, u) * kn(l, v) + v * jn(l, u) * kn(l - 1, v)) 103 | 104 | epsilon = 1.e-12 105 | 106 | Vc = get_mode_cutoffs(l+1, m)[-1] 107 | if V < Vc: 108 | lowbound = 0. 109 | else: 110 | lowbound = 1.-(Vc/V)**2 111 | Vcl = get_mode_cutoffs(l, m)[-1] 112 | if V < Vcl: 113 | return np.nan 114 | 115 | highbound = 1.-(Vcl/V)**2 116 | 117 | if np.isnan(lowbound): lowbound = 0. 118 | if np.isnan(highbound): highbound = 1. 119 | 120 | lowbound = np.max((lowbound-epsilon,0.+epsilon)) 121 | highbound = np.min((highbound+epsilon,1.)) 122 | b_opt = findBetween(solve_fn, lowbound, highbound, maxj=10, args=(V,)) 123 | 124 | return b_opt 125 | 126 | def lpfield(xg,yg,l,m,a,wl0,ncore,nclad,which = "cos"): 127 | '''calculate transverse field distribution of lp mode''' 128 | 129 | assert which in ("cos","sin"), "lp mode azimuthal component is either a cosine or sine, choose either 'cos' or 'sin'" 130 | 131 | V = get_V(2*np.pi/wl0,a,ncore,nclad) 132 | rs = np.sqrt(np.power(xg,2) + np.power(yg,2)) 133 | b = get_b(l,m,V) 134 | 135 | #print(np.sqrt(nclad**2+b*(ncore**2-nclad**2))) 136 | 137 | u = V*np.sqrt(1-b) 138 | v = V*np.sqrt(b) 139 | 140 | fieldout = np.zeros_like(rs,dtype = np.complex128) 141 | 142 | inmask = np.nonzero(rs<=a) 143 | outmask = np.nonzero(rs>a) 144 | 145 | fieldout[inmask] = jn(l,u*rs[inmask]/a) 146 | fieldout[outmask] = jn(l,u)/kn(l,v) * kn(l,v*rs[outmask]/a) 147 | 148 | #cosine/sine modulation 149 | phis = np.arctan2(yg,xg) 150 | if which == 'cos': 151 | fieldout *= np.cos(l*phis) 152 | else: 153 | fieldout *= np.sin(l*phis) 154 | 155 | return fieldout 156 | 157 | def get_IOR(wl): 158 | """ for fused silica """ 159 | wl2 = wl*wl 160 | return np.sqrt(0.6961663 * wl2 / (wl2 - 0.0684043**2) + 0.4079426 * wl2 / (wl2 - 0.1162414**2) + 0.8974794 * wl2 / (wl2 - 9.896161**2) + 1) 161 | 162 | 163 | def get_num_modes(k0,rcore,ncore,nclad): 164 | V = get_V(k0,rcore,ncore,nclad) 165 | modes = get_modes(V) 166 | num = 0 167 | for mode in modes: 168 | if mode[0]==0: 169 | num+=1 170 | else: 171 | num+=2 172 | return num 173 | 174 | def get_all_bs(l, V,bmax): 175 | 176 | def solve_fn(b,V): 177 | v = V*scimath.sqrt(b) 178 | u = V*scimath.sqrt(1.-b) 179 | 180 | if l == 0: 181 | return np.abs(u * jv(1, u) * kv(0, v) - v * jv(0, u) * kv(1, v)) 182 | else: 183 | return np.abs(u * jv(l - 1, u) * kv(l, v) + v * jv(l, u) * kv(l - 1, v)) 184 | 185 | from scipy.optimize import fsolve,minimize 186 | 187 | ret = minimize(solve_fn,[-131],args=(V,),bounds = [(None,bmax)]).x 188 | 189 | 190 | 191 | #return fsolve(solve_fn,-2,args=(V,)) 192 | 193 | 194 | if __name__ == "__main__": 195 | 196 | import matplotlib.pyplot as plt 197 | rcore = 21.8/2 198 | ncore = 1.4504 199 | nclad = 1.4504 - 5.5e-3 200 | 201 | import hcipy as hc 202 | 203 | V = get_V(2*np.pi,rcore,ncore,nclad) 204 | modes = get_modes(V) 205 | 206 | for mode in modes: 207 | 208 | xa = ya = np.linspace(-15,15,1000) 209 | xg , yg = np.meshgrid(xa,ya) 210 | 211 | if mode[0]==0: 212 | print(mode) 213 | lp= lpfield(xg,yg,mode[0],mode[1],rcore,1,ncore,nclad,'cos') 214 | lp /= np.max(lp) 215 | 216 | f = hc.Field(lp.flatten() , hc.make_pupil_grid((1000,1000),diameter=30) ) 217 | 218 | hc.imshow_field(f) 219 | plt.show() 220 | continue 221 | else: 222 | 223 | print(mode,"cos") 224 | lp= lpfield(xg,yg,mode[0],mode[1],rcore,1,ncore,nclad,'cos') 225 | lp /= np.max(lp) 226 | 227 | f = hc.Field(lp.flatten() , hc.make_pupil_grid((1000,1000),diameter=30) ) 228 | 229 | hc.imshow_field(f) 230 | plt.show() 231 | 232 | print(mode,'sin') 233 | lp= lpfield(xg,yg,mode[0],mode[1],rcore,1,ncore,nclad,'sin') 234 | lp /= np.max(lp) 235 | 236 | f = hc.Field(lp.flatten() , hc.make_pupil_grid((1000,1000),diameter=30) ) 237 | 238 | hc.imshow_field(f) 239 | plt.show() 240 | 241 | ''' 242 | rcore = 4 243 | wl = np.linspace(1.2,1.4,100) 244 | k0 = 2*np.pi/wl 245 | ncore = 1.4504 246 | nclad = 1.4504 - 5.5e-3 247 | 248 | modenumes = np.vectorize(get_num_modes)(2*np.pi/wl,rcore,ncore,nclad) 249 | import matplotlib.pyplot as plt 250 | plt.plot(wl,modenumes) 251 | plt.show() 252 | 253 | #k = 2*np.pi/0.98 254 | #print(get_NA(1.4504 + 0.0088,1.4504)) 255 | ''' 256 | 257 | """ 258 | rcore = 2.2 259 | NA = 0.16 260 | wl0 = 1.55 261 | ncore = 4 262 | nclad = np.sqrt(ncore*ncore-NA*NA) 263 | 264 | 265 | print(nclad,ncore) 266 | 267 | k0 = 2*np.pi/wl0 268 | 269 | print(get_MFD(k0,rcore,ncore,nclad)) 270 | """ 271 | """ 272 | import matplotlib.pyplot as plt 273 | 274 | wl = 1. 275 | k = 2*np.pi/wl 276 | ncore = 1.4504 277 | nclad = 1.4504 - 5.5e-3 278 | 279 | rcore = 12 280 | V = get_V(k,rcore,ncore,nclad) 281 | 282 | modes = get_modes(V) 283 | print(modes) 284 | xa = ya = np.linspace(-20,20,801) 285 | xg, yg = np.meshgrid(xa,ya) 286 | 287 | 288 | fig,axs = plt.subplots(7,6) 289 | 290 | for mode in modes: 291 | if mode[0] == 0: 292 | field = lpfield(xg,yg,mode[0],mode[1],rcore,wl,ncore,nclad) 293 | axs[0,2*mode[1]-2].imshow(np.real(field),vmin = -np.max(np.real(field)),vmax = np.max(np.real(field))) 294 | else: 295 | fieldcos = lpfield(xg,yg,mode[0],mode[1],rcore,wl,ncore,nclad,'cos') 296 | fieldsin = lpfield(xg,yg,mode[0],mode[1],rcore,wl,ncore,nclad,'sin') 297 | axs[mode[0],2*mode[1]-2].imshow(np.real(fieldcos),vmin = -np.max(np.real(fieldcos)),vmax = np.max(np.real(fieldcos))) 298 | axs[mode[0],2*mode[1]-1].imshow(np.real(fieldsin),vmin = -np.max(np.real(fieldsin)),vmax = np.max(np.real(fieldsin))) 299 | 300 | for _axs in axs: 301 | for ax in _axs: 302 | ax.set_frame_on(False) 303 | ax.axes.get_yaxis().set_visible(False) 304 | ax.axes.get_xaxis().set_visible(False) 305 | 306 | plt.subplots_adjust(hspace=0,wspace=0) 307 | 308 | plt.show() 309 | """ -------------------------------------------------------------------------------- /src/lightbeam/mesh.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from numpy import s_,arange,sqrt,power,complex64 as c64 3 | from scipy.interpolate import RectBivariateSpline 4 | import matplotlib.pyplot as plt 5 | from itertools import chain 6 | from bisect import bisect_left 7 | import numexpr as ne 8 | import math 9 | 10 | ## to do 11 | 12 | # figure out how to normalize ucrit 13 | # combine the different remeshing options into a single func with a switch argument 14 | # remeshing process is a little inefficient. not sure how to speed up though 15 | 16 | TOL=1e-12 17 | 18 | class RectMesh2D: 19 | ''' transverse adapative mesh class ''' 20 | 21 | def __init__(self,xw,yw,dx,dy,Nbc=4): 22 | 23 | self.max_iters = 6 24 | self.Nbc = Nbc 25 | self.dx0,self.dy0 = dx,dy 26 | 27 | self.xa = None 28 | self.ya = None 29 | 30 | self.xg = None 31 | self.yg = None 32 | 33 | self.ccel_ix = np.s_[Nbc+2:-Nbc-2] 34 | 35 | ###??? idk why these have to overlap but the pml doesn't work any other way 36 | self.cvert_ix = np.s_[Nbc:-Nbc] 37 | self.pvert_ix = np.hstack((arange(Nbc+1),arange(-Nbc-1,0))) 38 | 39 | self.reinit(xw,yw) 40 | self.update(self.rfacxa,self.rfacya) 41 | 42 | def reinit(self,xw,yw): 43 | self.xa_last,self.ya_last = self.xa,self.ya 44 | 45 | dx,dy = self.dx0,self.dy0 46 | Nbc = self.Nbc 47 | 48 | self.shape0_comp = (int(round(xw/dx)+1),int(round(yw/dy)+1)) 49 | xres,yres = self.shape0_comp[0] + 2*Nbc , self.shape0_comp[1] + 2*Nbc 50 | 51 | self.shape0 = (xres,yres) 52 | self.shape = (xres,yres) 53 | 54 | self.xw,self.yw = xw,yw 55 | 56 | #anchor point of the adaptive grid 57 | self.xm = -xw/2-Nbc*dx 58 | self.ym = -yw/2-Nbc*dy 59 | self.xM = xw/2+Nbc*dx 60 | self.yM = yw/2+Nbc*dy 61 | 62 | self.xa0 = np.linspace(-xw/2-Nbc*dx,xw/2+Nbc*dx,xres) 63 | self.ya0 = np.linspace(-yw/2-Nbc*dy,yw/2+Nbc*dy,yres) 64 | 65 | self.xix_base = np.arange(xres) 66 | self.yix_base = np.arange(yres) 67 | 68 | self.dxa = np.full(xres,dx) 69 | self.dya = np.full(yres,dy) 70 | 71 | self.xa = self.xa0 = np.linspace(-xw/2-Nbc*dx,xw/2+Nbc*dx,xres) 72 | self.ya = self.ya0 = np.linspace(-yw/2-Nbc*dy,yw/2+Nbc*dy,yres) 73 | 74 | self.pvert_xa = self.xa0[self.pvert_ix] 75 | self.pvert_ya = self.ya0[self.pvert_ix] 76 | 77 | self.rfacxa = self.rfacxa0 = np.full(xres-1,1) 78 | self.rfacya = self.rfacya0 = np.full(yres-1,1) 79 | 80 | def snapto(self,xw,yw): 81 | xwr = 2*math.ceil(xw/2/self.dx0) 82 | ywr = 2*math.ceil(yw/2/self.dy0) 83 | 84 | #round 85 | xw = xwr*self.dx0 86 | yw = ywr*self.dy0 87 | return xw, yw 88 | 89 | def dxa2xa(self,dxa): 90 | N = len(dxa) 91 | out = np.zeros(N+1) 92 | np.cumsum(dxa,out=out[1:]) 93 | return out + self.xm 94 | 95 | def update(self,rfacxa,rfacya): 96 | self.rfacxa = rfacxa 97 | self.rfacya = rfacya 98 | 99 | xix_base = np.empty(len(rfacxa)+1,dtype=int) 100 | yix_base = np.empty(len(rfacya)+1,dtype=int) 101 | 102 | xix_base[0] = 0 103 | yix_base[0] = 0 104 | 105 | xix_base[1:] = np.cumsum(rfacxa) 106 | yix_base[1:] = np.cumsum(rfacya) 107 | 108 | self.xix_base = xix_base[self.xix_base] 109 | self.yix_base = yix_base[self.yix_base] 110 | 111 | new_dxa = np.repeat(self.dxa[1:]/rfacxa,rfacxa) 112 | new_dya = np.repeat(self.dya[1:]/rfacya,rfacya) 113 | 114 | new_xa = self.dxa2xa(new_dxa) 115 | new_ya = self.dxa2xa(new_dya) 116 | 117 | self.xa = new_xa 118 | self.ya = new_ya 119 | 120 | rxa = np.empty_like(self.xa,dtype=float) 121 | rxa[1:-1] = new_dxa[1:]/new_dxa[:-1] 122 | rxa[0] = 1 123 | rxa[-1] = 1 124 | self.rxa = rxa 125 | 126 | rya = np.empty_like(self.ya,dtype=float) 127 | rya[1:-1] = new_dya[1:]/new_dya[:-1] 128 | rya[0] = 1 129 | rya[-1] = 1 130 | self.rya = rya 131 | 132 | self.dxa = np.empty_like(self.xa) 133 | self.dxa[1:] = new_dxa 134 | self.dxa[0] = self.dxa[1] 135 | 136 | self.dya = np.empty_like(self.ya) 137 | self.dya[1:] = new_dya 138 | self.dya[0] = self.dya[1] 139 | 140 | self.xres,self.yres = len(self.xa),len(self.ya) 141 | 142 | self.xg,self.yg = np.meshgrid(new_xa,new_ya,indexing='ij') 143 | 144 | #offset grids 145 | xhg = np.empty(( self.xg.shape[0] + 1 , self.xg.shape[1] )) 146 | yhg = np.empty(( self.yg.shape[0] , self.yg.shape[1] + 1 )) 147 | 148 | ne.evaluate("(a+b)/2",local_dict={"a":self.xg[1:],"b":self.xg[:-1]},out=xhg[1:-1]) 149 | ne.evaluate("(a+b)/2",local_dict={"a":self.yg[:,1:],"b":self.yg[:,:-1]},out=yhg[:,1:-1]) 150 | 151 | xhg[0] = self.xg[0] - self.dxa[0]*0.5 152 | xhg[-1] = self.xg[-1] + self.dxa[-1]*rxa[-1]*0.5 153 | 154 | yhg[:,0] = self.yg[:,0] - self.dya[0]*0.5 155 | yhg[:,-1] = self.yg[:,-1] + self.dya[-1]*self.rya[-1]*0.5 156 | 157 | self.xhg,self.yhg = xhg,yhg 158 | 159 | self.shape = (len(new_xa),len(new_ya)) 160 | 161 | def get_weights(self): 162 | xhg,yhg = self.xhg,self.yhg 163 | weights = ne.evaluate("(a-b)*(c-d)",local_dict={"a":xhg[1:],"b":xhg[:-1],"c":yhg[:,1:],"d":yhg[:,:-1]}) 164 | return weights 165 | 166 | def resample(self,u,xa=None,ya=None,newxa=None,newya=None): 167 | if xa is None or ya is None: 168 | out = RectBivariateSpline(self.xa_last,self.ya_last,u)(self.xa,self.ya) 169 | else: 170 | out = RectBivariateSpline(xa,ya,u)(newxa,newya) 171 | return out 172 | 173 | def resample_complex(self,u,xa=None,ya=None,newxa=None,newya=None): 174 | reals = np.real(u) 175 | imags = np.imag(u) 176 | reals = self.resample(reals,xa,ya,newxa,newya) 177 | imags = self.resample(imags,xa,ya,newxa,newya) 178 | return reals+1.j*imags 179 | 180 | def plot_mesh(self,reduce_by = 1,show=True): 181 | i=0 182 | for x in self.xa: 183 | if i%reduce_by == 0: 184 | plt.axhline(y=x,color='k',lw=0.5,alpha=0.5) 185 | i+=1 186 | i=0 187 | for y in self.ya: 188 | if i%reduce_by == 0: 189 | plt.axvline(x=y,color='k',lw=0.5,alpha=0.5) 190 | i+=1 191 | if show: 192 | plt.axis('equal') 193 | plt.show() 194 | 195 | def get_base_field(self,u): 196 | return u[self.xix_base].T[self.yix_base].T 197 | 198 | def _compute_refinement_factor(self,u0,crit_val): 199 | ''' given some electric field u0, compute values that corresponds to the degree of refinement required 200 | in the x and y subdivisions in the field. larger refinement factors should imply more grid refinement. 201 | crit_val determines the normalization of this refinement factor - this is left as an argument 202 | to allow the degree of mesh subdivision to be controlled. note that the output refinement factors, 203 | _rx and _ry, essentially will act as booleans. wherever these factors are larger than 1, the grid will 204 | marked for subdivision. 205 | ''' 206 | 207 | # the default scheme presented here sets the refinement factor to the geometric mean of field amplitude and 208 | # field second derivative magnitude. convergence testing shows that this metric leads to faster convergence 209 | # over using just field amplitude or second derivative magnitude alone. 210 | # evidence for this is entirely empirical and comes from testing that is not 100% comprehensive. 211 | 212 | ix = self.ccel_ix 213 | 214 | # x second derivative estimation 215 | xdif2 = np.empty_like(u0,dtype=np.complex128) 216 | xdif2[1:-1] = u0[2:]+u0[:-2] - 2*u0[1:-1] 217 | xdif2[0] = xdif2[-1] = 0 218 | 219 | # y second derivative estimation 220 | ydif2 = np.empty_like(u0,dtype=np.complex128) 221 | ydif2[:,1:-1] = u0[:,2:]+u0[:,:-2] - 2*u0[:,1:-1] 222 | 223 | ydif2[:,0] = ydif2[:,-1] = 0 224 | 225 | # field amps 226 | umaxx = np.sqrt(np.max(np.abs(u0),axis=1) * np.max(np.abs(xdif2),axis=1)) 227 | umaxx = 0.5*(umaxx[1:]+umaxx[:-1]) 228 | 229 | umaxy = np.sqrt(np.max(np.abs(u0),axis=0) * np.max(np.abs(ydif2),axis=0)) 230 | umaxy = 0.5*(umaxy[1:]+umaxy[:-1]) 231 | 232 | _rx = umaxx[ix]*self.dxa[1:][ix]/crit_val 233 | _ry = umaxy[ix]*self.dya[1:][ix]/crit_val 234 | 235 | return _rx,_ry 236 | 237 | 238 | def refine_by_two(self,u0,crit_val): 239 | ''' uses a hybrid approach where cells tagged are based on the product of field amplitude 240 | and second derivative magnitude ''' 241 | 242 | ix = self.ccel_ix 243 | 244 | _rx,_ry = self._compute_refinement_factor(u0,crit_val) 245 | 246 | rfacxa = np.full(u0.shape[0]-1,1,dtype=int) 247 | rfacya = np.full(u0.shape[1]-1,1,dtype=int) 248 | 249 | mask = (_rx>1) 250 | rfacxa[ix][mask] = 2 251 | 252 | mask = (_ry>1) 253 | rfacya[ix][mask] = 2 254 | 255 | xa_old = self.xa 256 | ya_old = self.ya 257 | 258 | self.update(rfacxa,rfacya) 259 | 260 | return self.resample_complex(u0,xa_old,ya_old,self.xa,self.ya) 261 | 262 | def refine_base(self,u0,ucrit): 263 | ''' iteratively apply refine_by_two to fully subdivide the simulation grid. ''' 264 | for i in range(self.max_iters): 265 | u0 = self.refine_by_two(u0,ucrit) 266 | 267 | class RectMesh3D: 268 | def __init__(self,xw,yw,zw,ds,dz,PML=4,xwfunc=None,ywfunc=None): 269 | '''base is a uniform mesh. can be refined''' 270 | self.xw,self.yw,self.zw = xw,yw,zw 271 | self.ds,self.dz = ds,dz 272 | self.xres,self.yres,self.zres = round(xw/ds)+1+2*PML, round(yw/ds)+1+2*PML, round(zw/dz)+1 273 | 274 | self.xa = np.linspace(-xw/2-PML*ds,xw/2+PML*ds,self.xres) 275 | self.ya = np.linspace(-xw/2-PML*ds,xw/2+PML*ds,self.yres) 276 | 277 | self.xg,self.yg = np.meshgrid(self.xa,self.ya,indexing='ij') 278 | 279 | self.shape=(self.zres,self.xres,self.yres) 280 | 281 | if xwfunc is None: 282 | xy_xw = xw 283 | else: 284 | xy_xw = 2*math.ceil(xwfunc(0)/2/ds)*ds 285 | 286 | if ywfunc is None: 287 | xy_yw = yw 288 | else: 289 | xy_yw = 2*math.ceil(ywfunc(0)/2/ds)*ds 290 | 291 | self.xy = RectMesh2D(xy_xw,xy_yw,ds,ds,PML) 292 | 293 | self.za = np.linspace(0,zw,self.zres) 294 | 295 | self.sigma_max = 5.+0.j #max (dimensionless) conductivity in PML layers 296 | self.PML = PML 297 | self.half_dz = dz/2. 298 | 299 | self.xwfunc = xwfunc 300 | self.ywfunc = ywfunc 301 | 302 | def get_loc(self ): 303 | 304 | xy = self.xy 305 | ix0 = bisect_left(self.xa,xy.xm-TOL) 306 | ix1 = bisect_left(self.xa,xy.xM-TOL) 307 | ix2 = bisect_left(self.ya,xy.ym-TOL) 308 | ix3 = bisect_left(self.ya,xy.yM-TOL) 309 | return ix0,ix1,ix2,ix3 310 | 311 | def sigmax(self,x): 312 | '''dimensionless, divided by e0 omega''' 313 | return np.where(np.abs(x)>self.xy.xw/2.,power((np.abs(x) - self.xy.xw/2)/(self.PML*self.xy.dx0),2.)*self.sigma_max,0.+0.j) 314 | 315 | def sigmay(self,y): 316 | '''dimensionless, divided by e0 omega''' 317 | return np.where(np.abs(y)>self.xy.yw/2.,power((np.abs(y) - self.xy.yw/2)/(self.PML*self.xy.dy0),2.)*self.sigma_max,0.+0.j) 318 | 319 | -------------------------------------------------------------------------------- /src/lightbeam/PIAA.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import matplotlib.pyplot as plt 3 | from scipy.interpolate import interp1d, UnivariateSpline 4 | import hcipy as hc 5 | from lightbeam import LPmodes 6 | from scipy.special import jn,kn 7 | from scipy.integrate import quad 8 | from scipy.optimize import brentq 9 | 10 | class PaddedWavefront(hc.Wavefront): 11 | def __init__(self, electric_field, wavelength=1, input_stokes_vector=None,pad_factor=1): 12 | shape = electric_field.shaped.shape 13 | assert shape[0] == shape[1] , "simulation region must be square!" 14 | assert shape[0]%2 == 0 , "simulation region must be an even number of pixels across!" 15 | 16 | res = shape[0] 17 | 18 | new_field = np.pad(electric_field.shaped,int(res*(pad_factor-1)/2)).flatten() 19 | 20 | self.radius = electric_field.grid[-1,0] 21 | self.total_radius = self.radius * pad_factor 22 | self.pad_factor = pad_factor 23 | self.beam_res = res 24 | self.res = pad_factor*res 25 | 26 | padded_electric_field = hc.Field(new_field,hc.make_pupil_grid(pad_factor*res,2*pad_factor*self.radius )) 27 | 28 | super().__init__(padded_electric_field, wavelength, input_stokes_vector) 29 | 30 | @staticmethod 31 | def remove_pad(wf,pad): 32 | res = wf.electric_field.shaped.shape[0] 33 | radius = wf.electric_field.grid[-1,0] 34 | beamres = int(res/pad) 35 | padpix = int( (pad-1)*beamres/2 ) 36 | field = wf.electric_field.shaped[padpix:-padpix,padpix:-padpix] 37 | grid = hc.make_pupil_grid(beamres,2*radius) 38 | return hc.Wavefront(hc.Field(field.flatten(),grid),wavelength = wf.wavelength) 39 | 40 | def make_remapping_gauss(res,obs,trunc=0.95,sigma='auto'): 41 | """make remapping arrays corresponding to gaussian INTENSITY output""" 42 | 43 | ra = np.linspace(0,1,res) #note that r=1 technically maps to infinity 44 | 45 | with np.errstate(divide='ignore'): 46 | ra_map = np.sqrt(-2 * np.log(1 - trunc*np.power(ra,2))) 47 | ra_map[0] = 0 48 | 49 | if sigma == 'auto': 50 | #normalize so pre and post remapping have same domain 51 | ra_map /= np.max(ra_map) 52 | else: 53 | ra_map *= sigma 54 | ra_obs = np.linspace(obs,1,res) 55 | 56 | return ra_obs, ra_map 57 | 58 | def make_remapping_gauss_annulus(res,obs,a,b,sigma='auto'): 59 | """make remapping arrays that map to gaussian 'annulus' with inner truncations a and outer truncation b (in units of sigma)""" 60 | ra = np.linspace(obs,1,res) 61 | out = np.sqrt( - 2*np.log( np.exp(-0.5*a*a) - (ra*ra-obs*obs)/(1-obs*obs) * (np.exp(-0.5*a*a) - np.exp(-0.5*b*b))) ) 62 | 63 | if sigma == 'auto': 64 | out /= (np.max(out)) 65 | else: 66 | out *= sigma 67 | 68 | return ra,out 69 | 70 | def make_remapping_gauss_annulus_2(res,obs,a,b): 71 | 72 | ra = np.linspace(obs,1,res) 73 | 74 | num = np.exp( a**2 + b**2 ) * (1 - obs*obs) 75 | 76 | den = np.exp( a**2 ) * (ra*ra - obs*obs) + np.exp(b**2) * (1 - ra*ra) 77 | out = np.sqrt(np.log(num/den)) 78 | out/=np.max(out) 79 | return ra,out 80 | 81 | def LP02_back(rcore,ncore,nclad,wl,dist): 82 | '''compute back propped LP02 mode. wl and rcore same units''' 83 | k0 = 2*np.pi/wl 84 | V = LPmodes.get_V(k0,rcore,ncore,nclad) 85 | b = LPmodes.get_b(0,2,V) 86 | u = V*np.sqrt(1-b) 87 | v = V*np.sqrt(b) 88 | 89 | def _inner_(rho): 90 | d = rho*2*np.pi*rcore/(wl*dist) 91 | term1 = (u * jn(1,u)*jn(0,d) - d * jn(0,u)*jn(1,d)) / (u**2-d**2) 92 | term2 = jn(0,u)/kn(0,u) * (d*kn(0,v)*jn(1,d) - v * kn(1,v)*jn(0,d)) / (v**2+d**2) 93 | 94 | return term1 - term2 95 | 96 | return _inner_ 97 | 98 | def LP0m_back(m,rcore,ncore,nclad,wl,dist,_norm=1): 99 | '''return a function for the backpropagated LP0m mode''' 100 | 101 | #convert to um, better for the integrator 102 | r_sc = rcore*1e6 103 | wl_sc = wl*1e6 104 | d_sc = dist*1e6 105 | 106 | k0 = 2*np.pi/wl_sc 107 | V = LPmodes.get_V(k0,r_sc,ncore,nclad) 108 | b = LPmodes.get_b(0,m,V) 109 | u = V*np.sqrt(1-b) 110 | v = V*np.sqrt(b) 111 | 112 | def _inner_(rho): 113 | d = rho*2*np.pi*r_sc/(wl_sc*d_sc) 114 | term1 = (u * jn(1,u)*jn(0,d) - d * jn(0,u)*jn(1,d)) / (u**2-d**2) 115 | term2 = jn(0,u)/kn(0,v) * (d*kn(0,v)*jn(1,d) - v * kn(1,v)*jn(0,d)) / (v**2+d**2) 116 | 117 | return term1 - term2 118 | 119 | # compute normalization prefactor 120 | 121 | def _integrand_(rho): 122 | return np.power(_inner_(rho),2)*rho 123 | 124 | norm_fac = np.sqrt( _norm/(quad(_integrand_,0,np.inf,limit=100)[0]*2*np.pi) ) 125 | 126 | def _inner2_(rho): 127 | d = rho*2*np.pi*r_sc/(wl_sc*d_sc) 128 | term1 = (u * jn(1,u)*jn(0,d) - d * jn(0,u)*jn(1,d)) / (u**2-d**2) 129 | term2 = jn(0,u)/kn(0,v) * (d*kn(0,v)*jn(1,d) - v * kn(1,v)*jn(0,d)) / (v**2+d**2) 130 | 131 | return norm_fac * (term1 - term2) 132 | 133 | return _inner2_ 134 | 135 | def make_remapping_LP0m(m,res,obs,beam_radius,rcore,ncore,nclad,wl,dist,inner_trunc=0,outer_trunc=None): 136 | if outer_trunc is None: 137 | outer_trunc = beam_radius 138 | 139 | back_lp0m_func = LP0m_back(m,rcore,ncore,nclad,wl,dist) 140 | 141 | r_in = beam_radius*obs 142 | inner_trunc = r_in * inner_trunc 143 | ra = np.linspace(r_in,beam_radius,res) 144 | 145 | def compute_encircled_power(_r,_norm=1): 146 | integrand = lambda r: r * np.power(back_lp0m_func(r*1e6),2) 147 | return _norm*quad(integrand,inner_trunc,_r)[0]*2*np.pi*1e12 148 | 149 | newnorm = 1/compute_encircled_power(outer_trunc) 150 | 151 | r_apos = [] 152 | for r in ra: 153 | if r == r_in: 154 | r_apos.append(inner_trunc) 155 | continue 156 | if r == beam_radius: 157 | r_apos.append(outer_trunc) 158 | continue 159 | 160 | f = lambda rho: compute_encircled_power(rho,newnorm) - (r*r-r_in*r_in)/(beam_radius*beam_radius-r_in*r_in) 161 | 162 | r_apo = brentq(f,inner_trunc,outer_trunc) 163 | r_apos.append(r_apo) 164 | 165 | _r = np.linspace(inner_trunc,outer_trunc,100) 166 | _p = np.vectorize(compute_encircled_power)(_r) 167 | 168 | #plt.title("encircled power") 169 | #plt.plot(_r,_p) 170 | 171 | #plt.show() 172 | 173 | return ra, np.array(r_apos) 174 | 175 | def make_PIAA_lenses(r1,r2,n1,n2,L): 176 | """ compute sag progiles for PIAA lenses to map r1 to r2 177 | 178 | args: 179 | r1,r2: remapping function (in array form) 180 | n1,n2: lens refractive index 181 | L: distance between lenses 182 | 183 | returns: 184 | z1,z2: height profiles of lenses 185 | """ 186 | z1 = np.zeros_like(r1) 187 | z2 = np.zeros_like(r2) 188 | 189 | z1[0] = 0 190 | z2[0] = L 191 | 192 | for i in range(len(r1)-1): 193 | _r1,_r2 = r1[i], r2[i] 194 | B = z2[i] - z1[i] 195 | A = _r1 - _r2 196 | slope1 = A / ( n1 * np.sqrt(A**2 + B**2) - B ) 197 | slope2 = A / ( n2 * np.sqrt(A**2 + B**2) - B ) 198 | z1[i+1] = -1*slope1 * (r1[i+1]-_r1) + z1[i] 199 | z2[i+1] = -1*slope2 * (r2[i+1]-_r2) + z2[i] 200 | 201 | return z1,z2 202 | 203 | def raytrace(r1,r2,z1,z2,n1,n2,L,skip=8): 204 | """use Snell's law to ray trace light paths between lenses""" 205 | 206 | #mirror lens profiles 207 | z1_p = np.concatenate( (z1[::-1][:-1],z1 ) ) 208 | z2_p = np.concatenate( (z2[::-1][:-1],z2 ) ) 209 | r1_p = np.concatenate( (-1*r1[::-1][:-1],r1 ) ) 210 | r2_p = np.concatenate( (-1*r2[::-1][:-1],r2 ) ) 211 | 212 | lens_thickness = 0.005 213 | 214 | #plt.axes().set_aspect('equal', 'datalim') 215 | 216 | #plot lens 1 217 | plt.plot(z1_p,r1_p,color='steelblue',lw=2) 218 | plt.plot((z1_p[-1]-lens_thickness,z1_p[-1]-lens_thickness), (r1_p[0],r1_p[-1]) ,color='steelblue',lw=2) 219 | 220 | plt.plot(z2_p,r2_p,color='steelblue',lw=2) 221 | plt.plot((z2_p[-1]+2*lens_thickness,z2_p[-1]+2*lens_thickness), (r2_p[0],r2_p[-1]) ,color='steelblue',lw=2) 222 | 223 | interpz1 = UnivariateSpline(r1,z1,s=0,k=1) 224 | interpz2 = UnivariateSpline(r2,z2,s=0,k=1) 225 | 226 | slope1 = interpz1.derivative(1) 227 | slope2 = interpz2.derivative(1) 228 | 229 | ray_rlaunch = r1[:-1] 230 | ray_zlaunch = z1[-1] - 2*lens_thickness 231 | 232 | for i in range(0,len(ray_rlaunch),skip): 233 | rayr = ray_rlaunch[i] 234 | _z1,_z2 = z1[i],z2[i] 235 | plt.plot( (ray_zlaunch,_z1 ), (rayr,rayr), color='k',lw=0.75) 236 | plt.plot( (ray_zlaunch,_z1 ), (-rayr,-rayr), color='k',lw=0.75) 237 | 238 | #first snell's law 239 | _slope = slope1(rayr) 240 | 241 | theta1 = np.arctan2(-1*_slope,1) 242 | theta2 = -np.arcsin(n1*np.sin(theta1)) + theta1 243 | 244 | _s = np.tan(theta2) 245 | 246 | ray_func1 = lambda z,z0,r0,: _s * (z-z0) + r0 247 | 248 | z_between = np.linspace(_z1,_z2,100) 249 | 250 | plt.plot(z_between,ray_func1(z_between,_z1,rayr),color='k',lw=0.75) 251 | plt.plot(z_between,-1*ray_func1(z_between,_z1,rayr),color='k',lw=0.75) 252 | 253 | #second snell's law 254 | 255 | new_rayr = ray_func1(_z2,_z1,rayr) 256 | _slope2 = slope2(new_rayr) 257 | 258 | theta_l2 = np.arctan2(-1*_slope2,1) 259 | 260 | theta3 = theta_l2 - theta2 261 | theta4 = np.arcsin(np.sin(theta3)/n2) - theta_l2 262 | 263 | _s2 = np.tan(theta4) 264 | 265 | ray_func2 = lambda z,z0,r0,: _s2 * (z-z0) + r0 266 | 267 | z_between_2 = np.linspace(_z2,z2[-1]+3*lens_thickness,20) 268 | 269 | plt.plot(z_between_2,ray_func2(z_between_2,_z2,new_rayr),color='k',lw=0.75) 270 | plt.plot(z_between_2,-1*ray_func2(z_between_2,_z2,new_rayr),color='k',lw=0.75) 271 | 272 | plt.show() 273 | 274 | def prop_through_lens(wf,z,n): 275 | wf = wf.copy() 276 | 277 | k = 2*np.pi / wf.wavelength 278 | #thickness = np.max(z) - np.min(z) 279 | phase_delay = k*(n-1)*z #k*n*z + k * (thickness - z) 280 | wf.electric_field *= np.exp(1j * phase_delay.flatten()) 281 | 282 | return wf 283 | 284 | def form_lens_height_maps(r1,r2,z1,z2,extent,res): 285 | """compute height of lens over uniform grid spanning simulation zone""" 286 | 287 | 288 | half_pixel = extent/res 289 | xa = ya = np.linspace(-extent+half_pixel,extent-half_pixel,res) # offset matches HCIPy's grids 290 | xg , yg = np.meshgrid(xa,ya) 291 | rg = np.sqrt(xg*xg+yg*yg) 292 | 293 | z1_func = UnivariateSpline(r1,z1,k=1,s=0,ext='const') 294 | z2_func = UnivariateSpline(r2,z2,k=1,s=0,ext='const') 295 | 296 | z1g = z1_func(rg) 297 | z1g-=np.min(z1g) 298 | 299 | z2g = z2_func(rg)*-1 300 | z2g-=np.min(z2g) 301 | 302 | return z1g,z2g 303 | 304 | def fresnel_apodizer(pupil_grid,radius,sep,r1,r2,z1,z2,IOR1,IOR2): 305 | # pupil_grid: regular Cartesian grid across beam 306 | # pad: padding factor for Fresnel prop section 307 | # r1,r2: aperture remapping arrays (r1,r2 are normalized) 308 | # z1,z2: 1D lens sag profiles 309 | # IOR, IOR2: refractive indices of first and second lens 310 | # res is resolution across the beam 311 | 312 | #radius = pupil_grid.x[-1] 313 | res = pupil_grid.shape[0] 314 | 315 | #r1 *= radius 316 | #r2 *= radius 317 | 318 | z1g , z2g = form_lens_height_maps(r1,r2,z1,z2,radius,res) 319 | 320 | prop = hc.AngularSpectrumPropagator(pupil_grid,sep) 321 | 322 | def _inner_(wf): 323 | wf = prop_through_lens(wf,z1g,IOR1) 324 | wf = prop(wf) 325 | wf = prop_through_lens(wf,z2g,IOR2) 326 | return wf 327 | 328 | def _inner_backwards(wf): 329 | wf = prop_through_lens(wf,-z2g,IOR2) 330 | wf = prop.backward(wf) 331 | wf = prop_through_lens(wf,-z1g,IOR1) 332 | return wf 333 | 334 | return _inner_,_inner_backwards 335 | 336 | if __name__ == "__main__": 337 | #plt.style.use("dark_background") 338 | 339 | rcore = 6.21 340 | ncore = 1.4504 341 | nclad = 1.4504 - 5.5e-3 342 | wl0 = 1.0 343 | sep = 25139 344 | beam_radius = 0.013/2*1e6 345 | piaa_sep = 0.12 346 | inner_trunc=0 347 | 348 | xa = ya = np.linspace(-6500,6500,1000) 349 | xg , yg = np.meshgrid(xa,ya) 350 | 351 | rg = np.sqrt(xg*xg+yg*yg) 352 | 353 | 354 | f = LP0m_back(2,rcore,ncore,nclad,wl0,sep) 355 | 356 | out = f(rg)*-1 357 | 358 | from misc import normalize 359 | 360 | out = normalize(out) 361 | 362 | plt.imshow( out) 363 | plt.show() 364 | 365 | r3,r4 = make_remapping_gauss_annulus(256,0.24,0,3) 366 | #r3,r4 = make_remapping_gauss_annulus_2(256,0.24,0,3/np.sqrt(2)) 367 | r3 *= beam_radius*1e-6 368 | r4 *= beam_radius*1e-6 369 | z3,z4 = make_PIAA_lenses(r3,r4,1.48,1.48,piaa_sep) 370 | 371 | raytrace(r3,r4,z3,z4,1.48,1.48,piaa_sep,8) 372 | 373 | 374 | r1,r2 = make_remapping_LP0m(2,200,0.24,beam_radius,rcore,ncore,nclad,wl0,sep,inner_trunc,None) 375 | 376 | r1/=1e6 377 | r2/=1e6 378 | 379 | plt.plot(np.zeros_like(r1),r1,ls='None',color='white',marker='.') 380 | plt.plot(np.ones_like(r2),r2,ls='None',color='white',marker='.') 381 | plt.show() 382 | 383 | z1,z2 = make_PIAA_lenses(r1,r2,1.48,1.48,piaa_sep) 384 | 385 | plt.plot(z1,r1) 386 | plt.plot(z2,r2) 387 | plt.axis('equal') 388 | plt.show() 389 | 390 | 391 | r1,r2 = make_remapping_gauss_annulus_2(100,0.24,0.24/np.sqrt(2),3/np.sqrt(2)) 392 | _r1,_r2 = make_remapping_gauss_annulus(100,0.24,0.24,3) 393 | 394 | plt.plot(np.zeros_like(r1),r1,marker='.',color='white',ls="None") 395 | plt.plot(np.ones_like(_r2),_r2,marker='.',color='steelblue',ls="None") 396 | plt.plot(np.ones_like(r2),r2,marker='.',color='white',ls="None") 397 | 398 | plt.ylim(0,1) 399 | plt.show() 400 | -------------------------------------------------------------------------------- /src/lightbeam/optics.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from numpy import logical_and as AND, logical_not as NOT 3 | from bisect import bisect_left,bisect_right 4 | import lightbeam.geom as geom 5 | from typing import List 6 | from lightbeam.mesh import RectMesh2D 7 | 8 | ### to do 9 | 10 | # some ideas 11 | # -currently have outer bbox to speed up raster. maybe an inner bbox will also speed things up? this would force fancy indexing though... 12 | # -but the boundary region between inner and outer bbox could also be split into four contiguous blocks 13 | # -extension to primitives with elliptical cross sections 14 | # -I don't think this is that hard. An ellipse is a stretched circle. So antialiasing the ellipse on a rectangular grid is the same 15 | # as antialiasing a circle on another differently stretched rectangular grid. 16 | # -but as of now, this is unnecessary 17 | 18 | class OpticPrim: 19 | '''base class for optical primitives (simple 3D shapes with a single IOR value)''' 20 | 21 | z_invariant = False 22 | 23 | def __init__(self,n): 24 | 25 | self.n = n 26 | self.n2 = n*n 27 | 28 | self.mask_saved = None 29 | 30 | # optionally, give prims a mesh to set the samplinig for IOR computations 31 | self.xymesh = None 32 | 33 | def _bbox(self,z): 34 | '''calculate the 2D bounding box of the primitive at given z. allows for faster IOR computation. Should be overwritten.''' 35 | return (-np.inf,np.inf,-np.inf,np.inf) 36 | 37 | def _contains(self,x,y,z): 38 | '''given coords, return whether or not those coords are inside the element. Should be overwritten.''' 39 | return np.full_like(x,False) 40 | 41 | def bbox_idx(self,z): 42 | '''get index slice corresponding to the primitives bbox, given an xg,yg coord grid''' 43 | m = self.xymesh 44 | xa,ya = m.xa,m.ya 45 | 46 | xmin,xmax,ymin,ymax = self._bbox(z) 47 | imin = max(bisect_left(xa,xmin)-1,0) 48 | imax = min(bisect_left(xa,xmax)+1,len(xa)) 49 | jmin = max(bisect_left(ya,ymin)-1,0) 50 | jmax = min(bisect_left(ya,ymax)+1,len(ya)) 51 | return np.s_[imin:imax,jmin:jmax], np.s_[imin:imax+1,jmin:jmax+1] 52 | 53 | def set_sampling(self,xymesh:RectMesh2D): 54 | self.xymesh = xymesh 55 | 56 | def set_IORsq(self,out,z,coeff=1): 57 | ''' replace values of out with IOR^2, given coordinate grids xg, yg, and z location. 58 | assumed primitive already has a set sampling grid''' 59 | 60 | if self.z_invariant and self.mask_saved is not None: 61 | mask = self.mask_saved 62 | else: 63 | bbox,bboxh = self.bbox_idx(z) 64 | mask = self._contains(self.xymesh.xg[bbox],self.xymesh.yg[bbox],z) 65 | 66 | if self.z_invariant and self.mask_saved is None: 67 | self.mask_saved = mask 68 | 69 | out[bbox][mask] = self.n2*coeff 70 | 71 | def get_boundary(self,z): 72 | ''' given z, get mask which will select pixels that lie on top of the primitive boundary 73 | you must set the sampling first before you call this!''' 74 | 75 | xhg,yhg = self.xymesh.xhg,self.xymesh.yhg 76 | maskh = self._contains(xhg,yhg,z) 77 | mask_r = (NOT ( AND(AND(maskh[1:,1:] == maskh[:-1,1:], maskh[1:,1:] == maskh[1:,:-1]),maskh[:-1,:-1]==maskh[:-1,1:]))) 78 | return mask_r 79 | 80 | class scaled_cyl(OpticPrim): 81 | ''' cylinder whose offset from origin and radius scale in the same way''' 82 | def __init__(self,xy,r,z_ex,n,nb,z_offset=0,scale_func=None,final_scale=1): 83 | ''' Initialize a scaled cylinder, where the cross-sectional geometry along the object's 84 | length is a scaled version of the initial geometry. 85 | 86 | Args: 87 | xy -- Initial location of the center of the cylinder at z=0. 88 | r -- initial cylinder radius 89 | z_ex -- cylinder length 90 | n -- refractive index of cylinder 91 | nb -- background index (required for anti-aliasing) 92 | 93 | z_offset -- offset that sets the z-coord for the cylinder's front 94 | scale_func -- optional custom function. Should take in z and return a scale value. 95 | set to None to use a linear scale function, where the scale factor 96 | of the back end is set by ... 97 | final_scale -- the scale of the final cross-section geometry, 98 | relative to the initial geoemtry. 99 | ''' 100 | 101 | super().__init__(n) 102 | self.p1 = p1 = [xy[0],xy[1],z_offset] 103 | self.p2 = [p1[0]*final_scale,p1[1]*final_scale,z_ex+z_offset] 104 | self.r = r 105 | self.rsq = r*r 106 | self.nb2 = nb*nb 107 | self.n2 = n*n 108 | self.z_ex = z_ex 109 | self.z_offset = z_offset 110 | 111 | self.AA = True 112 | 113 | def linear_func(_min,_max): 114 | slope = (_max - _min)/self.z_ex 115 | def _inner_(z): 116 | return slope*(z-self.z_offset) + _min 117 | return _inner_ 118 | 119 | if scale_func is None: 120 | scale_func = linear_func(1,final_scale) 121 | self.scale_func = scale_func 122 | self.xoffset_func = linear_func(p1[0],self.p2[0]) 123 | self.yoffset_func = linear_func(p1[1],self.p2[1]) 124 | 125 | def _contains(self,x,y,z): 126 | if not (self.z_offset <= z <= self.z_offset+self.z_ex): 127 | return False 128 | 129 | xdist = x - self.xoffset_func(z) 130 | ydist = y - self.yoffset_func(z) 131 | scale = self.scale_func(z) 132 | return (xdist*xdist + ydist*ydist <= scale*scale*self.rsq) 133 | 134 | def _bbox(self,z): 135 | xc = self.xoffset_func(z) 136 | yc = self.yoffset_func(z) 137 | scale = self.scale_func(z) 138 | xmax = xc+scale*self.r 139 | xmin = xc-scale*self.r 140 | ymax = yc+scale*self.r 141 | ymin = yc-scale*self.r 142 | return (xmin,xmax,ymin,ymax) 143 | 144 | def set_IORsq(self,out,z,coeff=1): 145 | '''overwrite base function to incorporate anti-aliasing and improve convergence''' 146 | if not (self.z_offset <= z <= self.z_offset+self.z_ex): 147 | return 148 | 149 | if not self.AA: 150 | super().set_IORsq(out,z,coeff) 151 | return 152 | 153 | center = (self.xoffset_func(z),self.yoffset_func(z)) 154 | scale = self.scale_func(z) 155 | bbox,bboxh = self.bbox_idx(z) 156 | xg = self.xymesh.xg[bbox] 157 | yg = self.xymesh.yg[bbox] 158 | 159 | xhg = self.xymesh.xhg[bboxh] 160 | yhg = self.xymesh.yhg[bboxh] 161 | 162 | m = self.xymesh 163 | rxg,ryg = np.meshgrid(m.rxa,m.rya,indexing='ij') 164 | dxg,dyg = np.meshgrid(m.dxa,m.dya,indexing='ij') 165 | 166 | geom.AA_circle_nonu(out,xg,yg,xhg,yhg,center,self.r*scale,self.nb2*coeff,self.n2*coeff,bbox,rxg,ryg,dxg,dyg) 167 | 168 | class OpticSys(OpticPrim): 169 | '''base class for optical systems, collections of primitives immersed in some background medium''' 170 | def __init__(self,elmnts:List[OpticPrim],nb): 171 | self.elmnts = elmnts 172 | self.nb = nb 173 | self.nb2 = nb*nb 174 | 175 | def _bbox(self,z): 176 | '''default behavior. won't work if the system has disjoint pieces''' 177 | if len(self.elmnts)==0: 178 | return super()._bbox(z) 179 | return self.elmnts[0]._bbox(z) 180 | 181 | def _contains(self,x,y,z): 182 | return self.elmnts[0]._contains(x,y,z) 183 | 184 | def set_sampling(self,xymesh:RectMesh2D): 185 | '''this function sets the spatial sampling for IOR computaitons''' 186 | super().set_sampling(xymesh) 187 | for elmnt in self.elmnts: 188 | elmnt.set_sampling(xymesh) 189 | 190 | def set_IORsq(self,out,z,xg=None,yg=None,coeff=1): 191 | '''replace values of out with IOR^2, given coordinate grids xg, yg, and z location.''' 192 | xg = self.xymesh.xg if xg is None else xg 193 | yg = self.xymesh.yg if yg is None else yg 194 | bbox,bboxh = self.bbox_idx(z) 195 | out[bbox] = self.nb2*coeff 196 | for elmnt in self.elmnts: 197 | elmnt.set_IORsq(out,z,coeff) 198 | 199 | class lant5(OpticSys): 200 | '''corrigan et al. 2018 style photonic lantern''' 201 | def __init__(self,rcore,rclad,rjack,ncore,nclad,njack,offset0,z_ex,scale_func=None,final_scale=1,nb=1): 202 | core0 = scaled_cyl([0,0],rcore,z_ex,ncore,nclad,scale_func=scale_func,final_scale=final_scale) 203 | core1 = scaled_cyl([offset0,0],rcore,z_ex,ncore,nclad,scale_func=scale_func,final_scale=final_scale) 204 | core2 = scaled_cyl([0,offset0],rcore,z_ex,ncore,nclad,scale_func=scale_func,final_scale=final_scale) 205 | core3 = scaled_cyl([-offset0,0],rcore,z_ex,ncore,nclad,scale_func=scale_func,final_scale=final_scale) 206 | core4 = scaled_cyl([0,-offset0],rcore,z_ex,ncore,nclad,scale_func=scale_func,final_scale=final_scale) 207 | clad = scaled_cyl([0,0],rclad,z_ex,nclad,njack,scale_func=scale_func,final_scale=final_scale) 208 | jack = scaled_cyl([0,0],rjack,z_ex,njack,nb,scale_func=scale_func,final_scale=final_scale) 209 | elmnts = [jack,clad,core4,core3,core2,core1,core0] 210 | 211 | super().__init__(elmnts,nb) 212 | 213 | class lant5big(OpticSys): 214 | '''corrigan et al. 2018 style photonic lantern except the jacket is infinite''' 215 | def __init__(self,rcore,rclad,ncore,nclad,njack,offset0,z_ex,scale_func=None,final_scale=1): 216 | core0 = scaled_cyl([0,0],rcore,z_ex,ncore,nclad,scale_func=scale_func,final_scale=final_scale) 217 | core1 = scaled_cyl([offset0,0],rcore,z_ex,ncore,nclad,scale_func=scale_func,final_scale=final_scale) 218 | core2 = scaled_cyl([0,offset0],rcore,z_ex,ncore,nclad,scale_func=scale_func,final_scale=final_scale) 219 | core3 = scaled_cyl([-offset0,0],rcore,z_ex,ncore,nclad,scale_func=scale_func,final_scale=final_scale) 220 | core4 = scaled_cyl([0,-offset0],rcore,z_ex,ncore,nclad,scale_func=scale_func,final_scale=final_scale) 221 | clad = scaled_cyl([0,0],rclad,z_ex,nclad,njack,scale_func=scale_func,final_scale=final_scale) 222 | elmnts = [clad,core4,core3,core2,core1,core0] 223 | 224 | super().__init__(elmnts,njack) 225 | 226 | class lant3big(OpticSys): 227 | '''3 port lantern, infinite jacket''' 228 | def __init__(self,rcore,rclad,ncore,nclad,njack,offset0,z_ex,z_offset=0,scale_func=None,final_scale=1): 229 | core0 = scaled_cyl([0,offset0],rcore,z_ex,ncore,nclad,z_offset,scale_func=scale_func,final_scale=final_scale) 230 | core1 = scaled_cyl([-np.sqrt(3)/2*offset0,-offset0/2],rcore,z_ex,ncore,nclad,z_offset,scale_func=scale_func,final_scale=final_scale) 231 | core2 = scaled_cyl([np.sqrt(3)/2*offset0,-offset0/2],rcore,z_ex,ncore,nclad,z_offset,scale_func=scale_func,final_scale=final_scale) 232 | clad = scaled_cyl([0,0],rclad,z_ex,nclad,njack,z_offset,scale_func=scale_func,final_scale=final_scale) 233 | elmnts = [clad,core2,core1,core0] 234 | 235 | super().__init__(elmnts,njack) 236 | 237 | self.init_core_locs = np.array([[0,offset0],[-np.sqrt(3)/2*offset0,-offset0/2],[np.sqrt(3)/2*offset0,-offset0/2]]) 238 | self.final_core_locs = self.init_core_locs*final_scale 239 | 240 | class lant3_ms(OpticSys): 241 | '''3 port lantern, infinite jacket, one core is bigger than the rest to accept LP01 mode.''' 242 | def __init__(self,rcore1,rcore2,rclad,ncore,nclad,njack,offset0,z_ex,z_offset=0,scale_func=None,final_scale=1): 243 | core0 = scaled_cyl([0,offset0],rcore1,z_ex,ncore,nclad,z_offset,scale_func=scale_func,final_scale=final_scale) 244 | core1 = scaled_cyl([-np.sqrt(3)/2*offset0,-offset0/2],rcore2,z_ex,ncore,nclad,z_offset,scale_func=scale_func,final_scale=final_scale) 245 | core2 = scaled_cyl([np.sqrt(3)/2*offset0,-offset0/2],rcore2,z_ex,ncore,nclad,z_offset,scale_func=scale_func,final_scale=final_scale) 246 | clad = scaled_cyl([0,0],rclad,z_ex,nclad,njack,z_offset,scale_func=scale_func,final_scale=final_scale) 247 | elmnts = [clad,core2,core1,core0] 248 | super().__init__(elmnts,njack) 249 | 250 | class lant6_saval(OpticSys): 251 | '''6 port lantern, mode-selective, based off sergio leon-saval's paper''' 252 | def __init__(self,rcore0,rcore1,rcore2,rcore3,rclad,ncore,nclad,njack,offset0,z_ex,z_offset=0,scale_func=None,final_scale=1): 253 | 254 | t = 2*np.pi/5 255 | core_locs = [[0,0]] 256 | for i in range(5): 257 | core_locs.append([offset0*np.cos(i*t),offset0*np.sin(i*t)]) 258 | self.core_locs = np.array(core_locs) 259 | core0 = scaled_cyl(core_locs[0],rcore0,z_ex,ncore,nclad,z_offset,scale_func=scale_func,final_scale=final_scale) 260 | core1 = scaled_cyl(core_locs[1],rcore1,z_ex,ncore,nclad,z_offset,scale_func=scale_func,final_scale=final_scale) 261 | core2 = scaled_cyl(core_locs[2],rcore1,z_ex,ncore,nclad,z_offset,scale_func=scale_func,final_scale=final_scale) 262 | core3 = scaled_cyl(core_locs[3],rcore2,z_ex,ncore,nclad,z_offset,scale_func=scale_func,final_scale=final_scale) 263 | core4 = scaled_cyl(core_locs[4],rcore2,z_ex,ncore,nclad,z_offset,scale_func=scale_func,final_scale=final_scale) 264 | core5 = scaled_cyl(core_locs[5],rcore3,z_ex,ncore,nclad,z_offset,scale_func=scale_func,final_scale=final_scale) 265 | 266 | clad = scaled_cyl([0,0],rclad,z_ex,nclad,njack,z_offset,scale_func=scale_func,final_scale=final_scale) 267 | elmnts = [clad,core5,core4,core3,core2,core1,core0] 268 | 269 | super().__init__(elmnts,njack) 270 | 271 | class lant19(OpticSys): 272 | '''19 port lantern, with cores hexagonally packed''' 273 | 274 | def __init__(self,rcore,rclad,ncore,nclad,njack,core_spacing,z_ex,z_offset=0,scale_func=None,final_scale=1): 275 | 276 | core_locs = self.get_19port_positions(core_spacing) 277 | self.core_locs = core_locs 278 | clad = scaled_cyl([0,0],rclad,z_ex,nclad,njack,z_offset,scale_func=scale_func,final_scale=final_scale) 279 | elmnts = [clad] 280 | 281 | for loc in core_locs: 282 | core = scaled_cyl(loc,rcore,z_ex,ncore,nclad,z_offset,scale_func=scale_func,final_scale=final_scale) 283 | elmnts.append(core) 284 | super().__init__(elmnts,njack) 285 | 286 | @staticmethod 287 | def get_19port_positions(core_spacing,plot=False): 288 | pos= [[0,0]] 289 | 290 | for i in range(6): 291 | xpos = core_spacing*np.cos(i*np.pi/3) 292 | ypos = core_spacing*np.sin(i*np.pi/3) 293 | pos.append([xpos,ypos]) 294 | 295 | startpos = np.array([2*core_spacing,0]) 296 | startang = 2*np.pi/3 297 | pos.append(startpos) 298 | for i in range(11): 299 | if i%2==0 and i!=0: 300 | startang += np.pi/3 301 | nextpos = startpos + np.array([core_spacing*np.cos(startang),core_spacing*np.sin(startang)]) 302 | pos.append(nextpos) 303 | startpos = nextpos 304 | 305 | pos = np.array(pos) 306 | if not plot: 307 | return pos 308 | 309 | import matplotlib.pyplot as plt 310 | 311 | for p in pos: 312 | plt.plot(*p,marker='.',ms=10,color='k') 313 | 314 | plt.axis('equal') 315 | plt.show() 316 | -------------------------------------------------------------------------------- /src/lightbeam/prop.py: -------------------------------------------------------------------------------- 1 | 2 | import numpy as np 3 | from numpy import exp,dot,full,cos,sin,real,imag,power,pi,log,sqrt,roll,linspace,arange,transpose,pad,complex128 as c128, float32 as f32, float64 as f64 4 | from numba import njit,jit,complex128 as nbc128, void 5 | import os 6 | os.environ['NUMEXPR_MAX_THREADS'] = '16' 7 | os.environ['NUMEXPR_NUM_THREADS'] = '8' 8 | import numexpr as ne 9 | from lightbeam.mesh import RectMesh3D,RectMesh2D 10 | import lightbeam.optics as optics 11 | from lightbeam.misc import timeit, overlap, normalize,printProgressBar, overlap_nonu, norm_nonu 12 | 13 | ### to do ### 14 | 15 | ## performance 16 | 17 | # maybe adaptive z stepping 18 | # get a better refinement criterion -- now weighting partially by 2nd deriv. still could use some work 19 | 20 | # more efficient ways to store arrays with many repeated values -- sparse data structure? 21 | 22 | ## readability 23 | 24 | def genc(shape): 25 | return np.empty(shape,dtype=c128,order='F') 26 | 27 | def genf(shape): 28 | return np.empty(shape,dtype=c128,order='F') 29 | 30 | @njit(void(nbc128[:,:],nbc128[:,:],nbc128[:,:],nbc128[:,:],nbc128[:,:],nbc128[:,:])) 31 | def tri_solve_vec(a,b,c,r,g,u): 32 | '''Apply Thomas' method for simultaneously solving a set of tridagonal systems. a, b, c, and r are matrices 33 | (N rows) where each column corresponds a separate system''' 34 | 35 | N = a.shape[0] 36 | beta = b[0] 37 | u[0] = r[0]/beta 38 | 39 | for j in range(1,N): 40 | g[j] = c[j-1]/beta 41 | beta = b[j] - a[j]*g[j] 42 | u[j] = (r[j] - a[j]*u[j-1])/beta 43 | 44 | for j in range(N-1): 45 | k = N-2-j 46 | u[k] = u[k] - g[k+1]*u[k+1] 47 | 48 | class Prop3D: 49 | '''beam propagator. employs finite-differences beam propagation with PML as the boundary condition. works on an adaptive mesh''' 50 | def __init__(self,wl0,_mesh:RectMesh3D,optical_system:optics.OpticSys,n0): 51 | 52 | xymesh = _mesh.xy 53 | 54 | self.wl0 = wl0 55 | self.k0 = k0 = 2.*pi/wl0 56 | self.k02 = k02 = k0*k0 57 | self._mesh = _mesh 58 | self.n0 = n0 59 | 60 | self.sig = sig = -2.j*k0*n0/_mesh.dz 61 | 62 | self.field = None 63 | 64 | self.optical_system = optical_system 65 | self.optical_system.set_sampling(xymesh) 66 | self.nb2 = nb2 = optical_system.nb2 67 | self.n02 = n02 = n0*n0 68 | 69 | ## things that will be set during computation 70 | 71 | self.xgrid_cor_facs = [[]]*3 72 | self.ygrid_cor_facs = [[]]*3 73 | 74 | self.xgrid_cor_mask = [] 75 | self.ygrid_cor_mask = [] 76 | 77 | ## precomputing some stuff ## 78 | 79 | Rx,Tupx,Tdox,Ry,Tupy,Tdoy = self.calculate_PML_mats() 80 | 81 | dx02 = _mesh.xy.dx0**2 82 | dy02 = _mesh.xy.dy0**2 83 | 84 | K = k02*(nb2-n02) 85 | n02 = power(n0,2) 86 | 87 | ## coeff matrices of tridiagonal system, updated periodically 88 | 89 | self._a0x = None 90 | self._b0x = None 91 | self._c0x = None 92 | self._a0y = None 93 | self._b0y = None 94 | self._c0y = None 95 | 96 | self.a0x_ = None 97 | self.b0x_ = None 98 | self.c0x_ = None 99 | self.a0y_ = None 100 | self.b0y_ = None 101 | self.c0y_ = None 102 | 103 | ## same as above but in PML zone 104 | 105 | self._apmlx = sig/12. - 0.5/dx02*Tdox - K/48. 106 | self._bpmlx = 5./6.*sig + Rx/dx02 - 5./24. * K 107 | self._cpmlx = sig/12. - 0.5/dx02*Tupx - K/48. 108 | 109 | self.apmlx_ = sig/12. + 0.5/dx02*Tdox + K/48. 110 | self.bpmlx_ = 5./6.*sig - Rx/dx02 + 5./24. * K 111 | self.cpmlx_ = sig/12. + 0.5/dx02*Tupx + K/48. 112 | 113 | self._apmly = sig/12. - 0.5/dy02*Tdoy - K/48. 114 | self._bpmly = 5./6.*sig + Ry/dy02 - 5./24. * K 115 | self._cpmly = sig/12. - 0.5/dy02*Tupy - K/48. 116 | 117 | self.apmly_ = sig/12. + 0.5/dy02*Tdoy + K/48. 118 | self.bpmly_ = 5./6.*sig - Ry/dy02 + 5./24. * K 119 | self.cpmly_ = sig/12. + 0.5/dy02*Tupy + K/48. 120 | 121 | self.half_dz = _mesh.dz/2. 122 | 123 | self.power = np.empty((_mesh.zres,)) 124 | self.totalpower = np.empty((_mesh.zres,)) 125 | 126 | def allocate_mats(self): 127 | sx,sy = self._mesh.xy.xg.shape,self._mesh.xy.yg.T.shape 128 | _trimatsx = (genc(sx),genc(sx),genc(sx)) 129 | _trimatsy = (genc(sy),genc(sy),genc(sy)) 130 | 131 | rmatx,rmaty = genc(sx),genc(sy) 132 | gx = genc(sx) 133 | gy = genc(sy) 134 | 135 | fill = self.nb2*self.k02 136 | 137 | IORsq__ = np.full(sx,fill,dtype=f64) 138 | _IORsq_ = np.full(sx,fill,dtype=f64) 139 | __IORsq = np.full(sx,fill,dtype=f64) 140 | 141 | return _trimatsx,rmatx,gx,_trimatsy,rmaty,gy,IORsq__,_IORsq_,__IORsq 142 | 143 | def check_z_inv(self): 144 | return self.optical_system.z_invariant 145 | 146 | def set_IORsq(self,out,z,xg=None,yg=None): 147 | #premultiply by k02 so we don't have to keep doing it later 148 | self.optical_system.set_IORsq(out,z,xg,yg,coeff=self.k02) 149 | 150 | def calculate_PML_mats(self): 151 | '''As per textbook , 152 | calculate the matrices R, T_j+1, and T_j-1 in the PML zone. We assume that the 153 | the PML's refractive index will be constant, equal to the background index. 154 | ''' 155 | 156 | m = self._mesh 157 | xy = m.xy 158 | 159 | xverts = xy.pvert_xa 160 | 161 | sdox = m.sigmax(xverts-xy.dx0) 162 | sx = m.sigmax(xverts) 163 | supx = m.sigmax(xverts+xy.dx0) 164 | 165 | yverts = xy.pvert_ya 166 | sdoy = m.sigmay(yverts-xy.dy0) 167 | sy = m.sigmay(yverts) 168 | supy = m.sigmay(yverts+xy.dy0) 169 | 170 | Qdox = 1./(1.+1.j*sdox*self.nb2) 171 | Qx = 1./(1.+1.j*sx*self.nb2) 172 | Qupx = 1./(1.+1.j*supx*self.nb2) 173 | 174 | Tupx = 0.5 * Qx * (Qx+Qupx) 175 | Tdox = 0.5 * Qx * (Qx+Qdox) 176 | Rx = 0.25 * Qx * (Qdox+2*Qx+Qupx) 177 | 178 | Qdoy = 1./(1.+1.j*sdoy*self.nb2) 179 | Qy = 1./(1.+1.j*sy*self.nb2) 180 | Qupy = 1./(1.+1.j*supy*self.nb2) 181 | 182 | Tupy= 0.5 * Qy * (Qy+Qupy) 183 | Tdoy = 0.5 * Qy * (Qy+Qdoy) 184 | Ry = 0.25 * Qy * (Qdoy+2*Qy+Qupy) 185 | 186 | return (Rx,Tupx,Tdox,Ry,Tupy,Tdoy) 187 | 188 | def update_grid_cor_facs(self,which='x'): 189 | xy = self._mesh.xy 190 | ix = xy.cvert_ix 191 | if which=='x': 192 | r = xy.rxa[ix] 193 | self.xgrid_cor_imask = np.where(r[1:-1]!=1)[0] 194 | else: 195 | r = xy.rya[ix] 196 | self.ygrid_cor_imask = np.where(r[1:-1]!=1)[0] 197 | r2 = r*r 198 | 199 | R1 = (r2 + r -1)/(6*r*(r+1)) 200 | R2 = (r2 + 3*r + 1)/(6*r) 201 | R3 = (-r2 + r + 1)/(6*(r+1)) 202 | 203 | ## alternative values from paper 204 | 205 | #R1 = (3*r2 - 3*r + 1)/ (6*r*(r+1)) 206 | #R2 = (-r2 + 7*r - 1)/(6*r) 207 | #R3 = (r2 - 3*r + 3)/(6*(r+1)) 208 | 209 | if which=='x': 210 | self.xgrid_cor_facs[0] = R1 211 | self.xgrid_cor_facs[1] = R2 212 | self.xgrid_cor_facs[2] = R3 213 | else: 214 | self.ygrid_cor_facs[0] = R1 215 | self.ygrid_cor_facs[1] = R2 216 | self.ygrid_cor_facs[2] = R3 217 | 218 | def precomp_trimats(self,which='x'): 219 | ix = self._mesh.xy.cvert_ix 220 | s = self.sig 221 | nu0 = -self.k02*self.n02 222 | 223 | eval1 = "s*r3 - 1/(r+1)/(d*d) - 0.25*r3*n" 224 | eval2 = "s*r2 + 1/r/(d*d) - 0.25*r2*n" 225 | eval3 = "s*r1 - 1/r/(r+1)/(d*d) - 0.25*r1*n" 226 | 227 | if which == 'x': 228 | R1,R2,R3 = self.xgrid_cor_facs 229 | r = self._mesh.xy.rxa[ix] 230 | dla = self._mesh.xy.dxa[ix] 231 | self._a0x = ne.evaluate(eval1,local_dict={"s":s,"r3":R3[1:,None],"r":r[1:,None],"d":dla[1:,None],"n":nu0}) 232 | self._b0x = ne.evaluate(eval2,local_dict={"s":s,"r2":R2[:,None],"r":r[:,None],"d":dla[:,None],"n":nu0}) 233 | self._c0x = ne.evaluate(eval3,local_dict={"s":s,"r1":R1[:-1,None],"r":r[:-1,None],"d":dla[:-1,None],"n":nu0}) 234 | else: 235 | R1,R2,R3 = self.ygrid_cor_facs 236 | r = self._mesh.xy.rya[ix] 237 | dla = self._mesh.xy.dya[ix] 238 | self._a0y = ne.evaluate(eval1,local_dict={"s":s,"r3":R3[1:,None],"r":r[1:,None],"d":dla[1:,None],"n":nu0}) 239 | self._b0y = ne.evaluate(eval2,local_dict={"s":s,"r2":R2[:,None],"r":r[:,None],"d":dla[:,None],"n":nu0}) 240 | self._c0y = ne.evaluate(eval3,local_dict={"s":s,"r1":R1[:-1,None],"r":r[:-1,None],"d":dla[:-1,None],"n":nu0}) 241 | 242 | def _trimats(self,out,IORsq,which='x'): 243 | ''' calculate the tridiagonal matrices in the computational zone ''' 244 | 245 | ix = self._mesh.xy.cvert_ix 246 | _IORsq = IORsq[ix] 247 | 248 | if which == 'x': 249 | R1,R2,R3 = self.xgrid_cor_facs 250 | r = self._mesh.xy.rxa[ix] 251 | dla = self._mesh.xy.dxa[ix] 252 | a,b,c = self._a0x,self._b0x,self._c0x 253 | 254 | else: 255 | R1,R2,R3 = self.ygrid_cor_facs 256 | r = self._mesh.xy.rya[ix] 257 | dla = self._mesh.xy.dya[ix] 258 | a,b,c = self._a0y,self._b0y,self._c0y 259 | 260 | _a,_b,_c = out 261 | 262 | s = self.sig 263 | 264 | eval1 = "a - 0.25*r3*n" 265 | eval2 = "b - 0.25*r2*n" 266 | eval3 = "c - 0.25*r1*n" 267 | 268 | ne.evaluate(eval1,local_dict={"a":a,"r3":R3[1:,None],"n":_IORsq[:-1]},out=_a[ix][1:]) 269 | ne.evaluate(eval2,local_dict={"b":b,"r2":R2[:,None],"n":_IORsq},out=_b[ix]) 270 | ne.evaluate(eval3,local_dict={"c":c,"r1":R1[:-1,None],"n":_IORsq[1:]},out=_c[ix][:-1]) 271 | 272 | _a[ix][0] = s*R3[0] - 1. / ((r[0]+1) * dla[0]*dla[0]) - 0.25*R3[0]*(_IORsq[0]-self.n02*self.k02) 273 | _c[ix][-1] = s*R1[-1] - 1/r[-1]/(r[-1]+1)/(dla[-1]*dla[-1]) - 0.25*R1[-1]*(_IORsq[-1]-self.n02*self.k02) 274 | 275 | def rmat_pmlcorrect(self,_rmat,u,which='x'): 276 | 277 | if which == 'x': 278 | apml,bpml,cpml = self.apmlx_,self.bpmlx_,self.cpmlx_ 279 | else: 280 | apml,bpml,cpml = self.apmly_,self.bpmly_,self.cpmly_ 281 | 282 | pix = self._mesh.xy.pvert_ix 283 | 284 | temp = np.empty_like(_rmat[pix]) 285 | 286 | temp[1:-1] = apml[1:-1,None]*u[pix-1][1:-1] + bpml[1:-1,None]*u[pix][1:-1] + cpml[1:-1,None]*u[pix+1][1:-1] 287 | 288 | temp[0] = bpml[0]*u[0] + cpml[0]*u[1] 289 | temp[-1] = apml[-1]*u[-2] + bpml[-1]*u[-1] 290 | 291 | _rmat[pix] = temp 292 | 293 | def rmat(self,_rmat,u,IORsq,which='x'): 294 | ix = self._mesh.xy.cvert_ix 295 | _IORsq = IORsq[ix] 296 | s = self.sig 297 | 298 | if which == 'x': 299 | R1,R2,R3 = self.xgrid_cor_facs 300 | dla = self._mesh.xy.dxa[ix] 301 | r = self._mesh.xy.rxa[ix] 302 | a,b,c = self.a0x_,self.b0x_,self.c0x_ 303 | else: 304 | R1,R2,R3 = self.ygrid_cor_facs 305 | dla = self._mesh.xy.dya[ix] 306 | r = self._mesh.xy.rya[ix] 307 | a,b,c = self.a0y_,self.b0y_,self.c0y_ 308 | 309 | N = self.n02*self.k02 310 | m = np.s_[1:-1,None] 311 | 312 | _dict = _dict = {"a":a,"b":b,"c":c,"u1":u[ix][:-2],"u2":u[ix][1:-1],"u3":u[ix][2:],"n1":_IORsq[:-2],"n2":_IORsq[1:-1],"n3":_IORsq[2:],"r3":R3[m],"r2":R2[m],"r1":R1[m] } 313 | _eval = "(a+0.25*r3*n1)*u1 + (b+0.25*r2*n2)*u2 + (c+0.25*r1*n3)*u3" 314 | 315 | ne.evaluate(_eval,local_dict=_dict,out=_rmat[ix][1:-1]) 316 | 317 | _rmat[ix][0] = (s*R2[0] - 1/(r[0]*dla[0]**2 ) + 0.25*R2[0]*(_IORsq[0]-N))*u[0] + (s*R1[0] + 1/r[0]/(r[0]+1)/dla[0]**2 + 0.25*R1[0] * (_IORsq[1]-N) )*u[1] 318 | _rmat[ix][-1] = (s*R3[-1] + 1. / ((r[-1]+1) * dla[-1]**2) + 0.25*R3[-1]*(_IORsq[-2]-N))*u[-2] + (s*R2[-1] - 1/(r[-1]*dla[-1]**2) + 0.25*R2[-1]*(_IORsq[-1]-N))*u[-1] 319 | 320 | def rmat_precomp(self,which='x'): 321 | ix = self._mesh.xy.cvert_ix 322 | s = self.sig 323 | n0 = -self.k02 * self.n02 324 | m = np.s_[1:-1,None] 325 | 326 | eval1="(s*r3+1/(r+1)/(d*d)+0.25*r3*n)" 327 | eval2="(s*r2-1/r/(d*d)+0.25*r2*n)" 328 | eval3="(s*r1+1/r/(r+1)/(d*d) + 0.25*r1*n)" 329 | 330 | if which == 'x': 331 | R1,R2,R3 = self.xgrid_cor_facs 332 | r = self._mesh.xy.rxa[ix] 333 | dla = self._mesh.xy.dxa[ix] 334 | 335 | _dict = {"s":s,"r3":R3[m],"r":r[m],"d":dla[m],"n":n0,"r2":R2[m],"r1":R1[m]} 336 | self.a0x_ = ne.evaluate(eval1,local_dict=_dict) 337 | self.b0x_ = ne.evaluate(eval2,local_dict=_dict) 338 | self.c0x_ = ne.evaluate(eval3,local_dict=_dict) 339 | else: 340 | R1,R2,R3 = self.ygrid_cor_facs 341 | r = self._mesh.xy.rya[ix] 342 | dla = self._mesh.xy.dya[ix] 343 | 344 | _dict = {"s":s,"r3":R3[m],"r":r[m],"d":dla[m],"n":n0,"r2":R2[m],"r1":R1[m]} 345 | self.a0y_ = ne.evaluate(eval1,local_dict=_dict) 346 | self.b0y_ = ne.evaluate(eval2,local_dict=_dict) 347 | self.c0y_ = ne.evaluate(eval3,local_dict=_dict) 348 | 349 | def _pmlcorrect(self,_trimats,which='x'): 350 | ix = self._mesh.xy.pvert_ix 351 | _a,_b,_c = _trimats 352 | 353 | if which=='x': 354 | _a[ix] = self._apmlx[:,None] 355 | _b[ix] = self._bpmlx[:,None] 356 | _c[ix] = self._cpmlx[:,None] 357 | else: 358 | _a[ix] = self._apmly[:,None] 359 | _b[ix] = self._bpmly[:,None] 360 | _c[ix] = self._cpmly[:,None] 361 | 362 | @timeit 363 | def prop2end(self,_u,monitor_func=None,ref_val=5.e-6,remesh_every=20,dynamic_n0 = False,fplanewidth=0,writeto=None,xyslice=None,zslice=None): 364 | 365 | _mesh = self._mesh 366 | PML = _mesh.PML 367 | 368 | if writeto is not None: 369 | # then we need to save some field information 370 | if xyslice is None and zslice is None: 371 | # save everything 372 | za_keep = _mesh.za 373 | shape = (len(za_keep),*_mesh.xg.shape) 374 | else: 375 | # save a slice 376 | za_keep = _mesh.za[zslice] 377 | shape = (len(za_keep),*_mesh.xg[xyslice].shape) 378 | 379 | self.field = np.zeros(shape,dtype=c128) 380 | 381 | #pull xy mesh 382 | xy = _mesh.xy 383 | dx,dy = xy.dx0,xy.dy0 384 | 385 | if fplanewidth == 0: 386 | xa_in = np.linspace(-_mesh.xw/2,_mesh.xw/2,xy.shape0_comp[0]) 387 | ya_in = np.linspace(-_mesh.yw/2,_mesh.yw/2,xy.shape0_comp[1]) 388 | else: 389 | xa_in = np.linspace(-fplanewidth/2,fplanewidth/2,xy.shape0_comp[0]) 390 | ya_in = np.linspace(-fplanewidth/2,fplanewidth/2,xy.shape0_comp[1]) 391 | 392 | dx0 = xa_in[1]-xa_in[0] 393 | dy0 = ya_in[1]-ya_in[0] 394 | 395 | # u can either be a field or a function that generates a field. 396 | # the latter option allows for coarse base grids to be used 397 | # without being penalized by forcing the use of a low resolution 398 | # launch field 399 | 400 | if type(_u) is np.ndarray: 401 | 402 | _power = overlap(_u,_u) 403 | print('input power: ',_power) 404 | 405 | # normalize the field, preserving the input power. accounts for grid resolution 406 | normalize(_u,weight=dx0*dy0,normval=_power) 407 | 408 | #resample the field onto the smaller xy mesh (in the smaller mesh's computation zone) 409 | u0 = xy.resample_complex(_u,xa_in,ya_in,xy.xa[PML:-PML],xy.ya[PML:-PML]) 410 | 411 | _power2 = overlap(u0,u0,dx*dy) 412 | 413 | #now we pad w/ zeros to extend it into the PML zone 414 | u0 = np.pad(u0,((PML,PML),(PML,PML))) 415 | 416 | #initial mesh refinement 417 | xy.refine_base(u0,ref_val) 418 | 419 | weights = xy.get_weights() 420 | 421 | #now resample the field onto the smaller *non-uniform* xy mesh 422 | u = xy.resample_complex(_u,xa_in,ya_in,xy.xa[PML:-PML],xy.ya[PML:-PML]) 423 | u = np.pad(u,((PML,PML),(PML,PML))) 424 | 425 | #do another norm to correct for the slight power change you get when resampling. I measure 0.1% change for psflo. should check again 426 | norm_nonu(u,weights,_power2) 427 | 428 | elif callable(_u): 429 | # must be of the form u(x,y) 430 | u0 = _u(xy.xg,xy.yg) 431 | _power = overlap(u0,u0) 432 | print('input power: ',_power) 433 | 434 | # normalize the field, preserving the input power. accounts for grid resolution 435 | normalize(u0,weight=dx0*dy0,normval=_power) 436 | 437 | # do an initial mesh refinement 438 | xy.refine_base(u0,ref_val) 439 | 440 | # compute the field on the nonuniform grid 441 | u = norm_nonu(_u(xy.xg,xy.yg),xy.get_weights(),_power) 442 | 443 | else: 444 | raise Exception("unsupported type for argument u in prop2end()") 445 | 446 | counter = 0 447 | total_iters = self._mesh.zres 448 | 449 | print("propagating field...") 450 | 451 | __z = 0 452 | z__ = 0 453 | 454 | #step 0 setup 455 | 456 | self.update_grid_cor_facs('x') 457 | self.update_grid_cor_facs('y') 458 | 459 | # initial array allocation 460 | _trimatsx,rmatx,gx,_trimatsy,rmaty,gy,IORsq__,_IORsq_,__IORsq = self.allocate_mats() 461 | 462 | self.precomp_trimats('x') 463 | self.precomp_trimats('y') 464 | 465 | self.rmat_precomp('x') 466 | self.rmat_precomp('y') 467 | 468 | self._pmlcorrect(_trimatsx,'x') 469 | self._pmlcorrect(_trimatsy,'y') 470 | 471 | #get the current IOR dist 472 | self.set_IORsq(IORsq__,z__) 473 | 474 | #plt.figure(frameon=False) 475 | #plt.imshow(xy.get_base_field(IORsq__)) 476 | #plt.show() 477 | 478 | print("initial shape: ",xy.shape) 479 | for i in range(total_iters): 480 | if i%20 == 0: 481 | printProgressBar(i,total_iters-1) 482 | u0 = xy.get_base_field(u) 483 | u0c = np.conj(u0) 484 | weights = xy.get_weights() 485 | 486 | ## Total power monitor ## 487 | self.totalpower[i] = overlap_nonu(u,u,weights) 488 | #print(self.totalpower[i]) 489 | 490 | ## Other monitors ## 491 | if monitor_func is not None: 492 | monitor_field = norm_nonu(monitor_func(xy.xg,xy.yg),weights) 493 | self.power[i] = power(overlap_nonu(u,monitor_field,weights),2) 494 | 495 | _z_ = z__ + _mesh.half_dz 496 | __z = z__ + _mesh.dz 497 | 498 | if self.field is not None and counter < len(za_keep) and __z == za_keep[counter]: 499 | # record the field 500 | self.field[counter] = u0[xyslice] 501 | counter += 1 502 | 503 | #avoid remeshing on step 0 504 | if (i+1)%remesh_every== 0: 505 | 506 | ## update the effective index 507 | if dynamic_n0: 508 | #update the effective index 509 | base = xy.get_base_field(IORsq__) 510 | self.n02 = xy.dx0*xy.dy0*np.real(np.sum(u0c*u0*base))/self.k02 511 | 512 | oldxm,oldxM = xy.xm,xy.xM 513 | oldym,oldyM = xy.ym,xy.yM 514 | 515 | oldxw,oldyw = xy.xw,xy.yw 516 | 517 | new_xw,new_yw = oldxw,oldyw 518 | #expand the grid if necessary 519 | if _mesh.xwfunc is not None: 520 | new_xw = _mesh.xwfunc(__z) 521 | if _mesh.ywfunc is not None: 522 | new_yw = _mesh.ywfunc(__z) 523 | 524 | new_xw, new_yw = xy.snapto(new_xw,new_yw) 525 | 526 | xy.reinit(new_xw,new_yw) #set grid back to base res with new dims 527 | 528 | if (xy.xw > oldxw or xy.yw > oldyw): 529 | #now we need to pad u,u0 with zeros to make sure it matches the new space 530 | xpad = int((xy.shape0[0]-u0.shape[0])/2) 531 | ypad = int((xy.shape0[1]-u0.shape[1])/2) 532 | 533 | u = np.pad(u,((xpad,xpad),(ypad,ypad))) 534 | u0 = np.pad(u0,((xpad,xpad),(ypad,ypad))) 535 | 536 | #pad coord arrays to do interpolation 537 | xy.xa_last = np.hstack( ( np.linspace(xy.xm,oldxm-dx,xpad) , xy.xa_last , np.linspace(oldxM + dx, xy.xM,xpad) ) ) 538 | xy.ya_last = np.hstack( ( np.linspace(xy.ym,oldym-dy,ypad) , xy.ya_last , np.linspace(oldyM + dy, xy.yM,ypad) ) ) 539 | 540 | #subdivide into nonuniform grid 541 | xy.refine_base(u0,ref_val) 542 | 543 | #interp the field to the new grid 544 | u = xy.resample_complex(u) 545 | 546 | #give the grid to the optical sys obj so it can compute IORs 547 | self.optical_system.set_sampling(xy) 548 | 549 | #compute nonuniform grid correction factors R_i 550 | self.update_grid_cor_facs('x') 551 | self.update_grid_cor_facs('y') 552 | 553 | # grid size has changed, so now we need to reallocate arrays for at least the next remesh_period iters 554 | _trimatsx,rmatx,gx,_trimatsy,rmaty,gy,IORsq__,_IORsq_,__IORsq = self.allocate_mats() 555 | 556 | #get the current IOR dist 557 | self.set_IORsq(IORsq__,z__) 558 | 559 | #precompute things that will be reused 560 | self.precomp_trimats('x') 561 | self.precomp_trimats('y') 562 | 563 | self.rmat_precomp('x') 564 | self.rmat_precomp('y') 565 | 566 | self._pmlcorrect(_trimatsx,'x') 567 | self._pmlcorrect(_trimatsy,'y') 568 | 569 | self.set_IORsq(_IORsq_,_z_,) 570 | self.set_IORsq(__IORsq,__z) 571 | 572 | self.rmat(rmatx,u,IORsq__,'x') 573 | self.rmat_pmlcorrect(rmatx,u,'x') 574 | 575 | self._trimats(_trimatsx,_IORsq_,'x') 576 | self._trimats(_trimatsy,__IORsq.T,'y') 577 | 578 | tri_solve_vec(_trimatsx[0],_trimatsx[1],_trimatsx[2],rmatx,gx,u) 579 | 580 | self.rmat(rmaty,u.T,_IORsq_.T,'y') 581 | self.rmat_pmlcorrect(rmaty,u.T,'y') 582 | 583 | tri_solve_vec(_trimatsy[0],_trimatsy[1],_trimatsy[2],rmaty,gy,u.T) 584 | 585 | z__ = __z 586 | if (i+2)%remesh_every != 0: 587 | IORsq__[:,:] = __IORsq 588 | 589 | print("final total power",self.totalpower[-1]) 590 | 591 | if writeto: 592 | np.save(writeto,self.field) 593 | return u,u0 594 | 595 | @timeit 596 | def prop2end_uniform(self,u,xyslice=None,zslice=None,u1_func=None,writeto=None,dynamic_n0 = False,fplanewidth=0): 597 | _mesh = self._mesh 598 | PML = _mesh.PML 599 | 600 | if not (xyslice is None and zslice is None): 601 | za_keep = _mesh.za[zslice] 602 | if type(za_keep) == np.ndarray: 603 | minz, maxz = za_keep[0],za_keep[-1] 604 | shape = (len(za_keep),*_mesh.xg[xyslice].shape) 605 | else: 606 | raise Exception('uhh not implemented') 607 | 608 | self.field = np.zeros(shape,dtype=c128) 609 | 610 | if fplanewidth == 0: 611 | xa_in = np.linspace(-_mesh.xw/2,_mesh.xw/2,u.shape[0]) 612 | ya_in = np.linspace(-_mesh.yw/2,_mesh.yw/2,u.shape[1]) 613 | else: 614 | xa_in = np.linspace(-fplanewidth/2,fplanewidth/2,u.shape[0]) 615 | ya_in = np.linspace(-fplanewidth/2,fplanewidth/2,u.shape[1]) 616 | 617 | dx0 = xa_in[1]-xa_in[0] 618 | dy0 = ya_in[1]-ya_in[0] 619 | 620 | _power = overlap(u,u) 621 | print('input power: ',_power) 622 | 623 | # normalize the field, preserving the input power. accounts for grid resolution 624 | normalize(u,weight=dx0*dy0,normval=_power) 625 | 626 | __z = 0 627 | 628 | #pull xy mesh 629 | xy = _mesh.xy 630 | dx,dy = xy.dx0,xy.dy0 631 | 632 | #resample the field onto the smaller xy mesh (in the smaller mesh's computation zone) 633 | u0 = xy.resample_complex(u,xa_in,ya_in,xy.xa[PML:-PML],xy.ya[PML:-PML]) 634 | 635 | #now we pad w/ zeros to extend it into the PML zone 636 | u0 = np.pad(u0,((PML,PML),(PML,PML))) 637 | 638 | counter = 0 639 | total_iters = self._mesh.zres 640 | 641 | print("propagating field...") 642 | 643 | z__ = 0 644 | 645 | #step 0 setup 646 | 647 | self.update_grid_cor_facs('x') 648 | self.update_grid_cor_facs('y') 649 | 650 | # initial array allocation 651 | _trimatsx,rmatx,gx,_trimatsy,rmaty,gy,IORsq__,_IORsq_,__IORsq = self.allocate_mats() 652 | 653 | self.precomp_trimats('x') 654 | self.precomp_trimats('y') 655 | 656 | self.rmat_precomp('x') 657 | self.rmat_precomp('y') 658 | 659 | self._pmlcorrect(_trimatsx,'x') 660 | self._pmlcorrect(_trimatsy,'y') 661 | 662 | #get the current IOR dist 663 | self.set_IORsq(IORsq__,z__) 664 | 665 | weights = xy.get_weights() 666 | 667 | print("initial shape: ",xy.shape) 668 | for i in range(total_iters): 669 | if i%20 == 0: 670 | printProgressBar(i,total_iters-1) 671 | 672 | ## Total power monitor ## 673 | self.totalpower[i] = overlap_nonu(u0,u0,weights) 674 | 675 | ## Other monitors ## 676 | if u1_func is not None: 677 | lp = norm_nonu(u1_func(xy.xg,xy.yg),weights) 678 | self.power[i] = power(overlap_nonu(u0,lp,weights),2) 679 | 680 | _z_ = z__ + _mesh.half_dz 681 | __z = z__ + _mesh.dz 682 | 683 | if self.field is not None and (minz<=__z<=maxz): 684 | ix0,ix1,ix2,ix3 = _mesh.get_loc() 685 | mid = int(u0.shape[1]/2) 686 | 687 | self.field[counter][ix0:ix1+1] = u0[:,mid] ## FIX ## 688 | counter+=1 689 | 690 | self.set_IORsq(_IORsq_,_z_,) 691 | self.set_IORsq(__IORsq,__z) 692 | 693 | self.rmat(rmatx,u0,IORsq__,'x') 694 | self.rmat_pmlcorrect(rmatx,u0,'x') 695 | 696 | self._trimats(_trimatsx,_IORsq_,'x') 697 | self._trimats(_trimatsy,__IORsq.T,'y') 698 | 699 | tri_solve_vec(_trimatsx[0],_trimatsx[1],_trimatsx[2],rmatx,gx,u0) 700 | 701 | self.rmat(rmaty,u0.T,_IORsq_.T,'y') 702 | self.rmat_pmlcorrect(rmaty,u0.T,'y') 703 | 704 | tri_solve_vec(_trimatsy[0],_trimatsy[1],_trimatsy[2],rmaty,gy,u0.T) 705 | 706 | z__ = __z 707 | IORsq__[:,:] = __IORsq 708 | 709 | print("final total power",self.totalpower[-1]) 710 | 711 | if writeto: 712 | np.save(writeto,self.field) 713 | return u0 714 | --------------------------------------------------------------------------------