Source code for finesse.knm.maps

# cython: binding=True

from libc.string cimport memcpy
import numpy as np
cimport numpy as np
import cython
import types

from .integrators import (
    outer_conj_product,
    outer_conj_product_2,
    compute_map_scattering_coeffs_riemann_optimised,
)
from .. import BeamParam
from ..env import warn
from ..utilities.maps import (
    overlap_piston_coefficient,
    overlap_curvature_coefficients,
    overlap_tilt_coefficients,
    rms
)
from ..cymath.homs import HGModes
from ..cymath.homs cimport unm_ws_recache_from_bp, unm_factor_store_init
from ..cymath.complex cimport complex_t
from ..cymath.gaussbeam cimport bp_beamsize
from cpython.ref cimport Py_XINCREF, Py_XDECREF
from libc.stdlib cimport free, calloc
from finesse.exceptions import FinesseException

[docs]cdef class Map: """A `Map` is an object that stores an optical path difference (OPD) and amplitude variations over some 2D surface. This represents a spatial distortion to an optical field and is used to model features such as surface defects, thermal lenses, and apertures. The complex value of the map is described: z(x,y) = amplitude(x,y) * exp[-1j * k * factor * opd(x,y)] Where k is the wavenumber for the wavelength being used, and factor is some
arbitrary scaling factor. This scaling is 2 for reflections, and for transmission will be some difference in refractive index between the two surfaces. Multiple `z` arrays can be generated from the same maps. `z` arrays are generated using the `get_z` method. Multiple options exist for automatic removal of tilts, curvatures, and astigmatisms in maps too. When these options are defined and new map data is generated by calling `get_z` these terms will be removed from the OPD data. This allows changing maps defined with functions to constantly have these terms being removed during a simulation, simplifying their use. When automatic removal is used the terms are weighted with the current complex beam parameter spot size as a weight. Parameters ---------- x : double[::1] 1D uniformally sample horizontal x array of size Nx y : double[::1] 1D uniformally sample vertical y array of size Ny opd : [double[:, ::1] | callable], optional 2D array of size [Ny, Nx] which describes the optical path difference over the grid defined by x and y. Units of meters. If a callable is provided it should be of the form `f(smap, Model)` argument and return a OPD array when called. `smap` is the current map object and `model` is the current model this Map's element is a part of. This allows you to access other elements and information if needed. amplitude : [double[:, ::1] | callable], optional 2D array of size [Ny, Nx] which describes a amplitude over the grid defined by x and y. Units of meters. If a callable is provided it should be of the form `f(smap, Model)` argument and return an amplitude array when called. `smap` is the current map object and `model` is the current model this Map's element is a part of. This allows you to access other elements and information if needed. auto_remove_tilts : bool When True any tilts present in the map will be removed when used during simulations. Can only be used if a callable is specified for ``opd`` then the tilts will be removed after each call to get new map data from the function. auto_remove_curvatures : bool When True any curvatures present in the map will be removed when used during simulations. Can only be used if callable is specified for ``opd`` then the tilts will be removed after each call to get new map data from the function. This will remove the average of the x and y curvatures, so will still leave astigmatisms present. auto_remove_astigmatism : bool When True any astigmatism present in the map will be removed when used during simulations. Can only be used if a callable is specified for ``opd`` then the tilts will be removed after each call to get new map data from the function. This will remove astigmatic curvatures in both x and y, so is similar to using `auto_remove_curvatures`. is_focusing_element : bool When True this map is representing a focusing element like a curved mirror or lens. In such cases the large scale curvature has not been removed from the map. When set to `True` the mode projection will be done using the map integration and q_in != q_out. Therefore this map must include the focusing OPD to correctly project q_in -> q_out to be mode matched. If False, the input and output q value is the same, and the mode projection is done using the analytic Bayer-Helms solution to the overlap integrals. put_focal_length : :class:`Parameter` A lens element added to the model to put any removed focal length in the map OPD into to. This is only updated when the :meth:`remove_curvatures` is called. This can happen at the user request or automatically if `auto_remove_curvatures` or `auto_remove_astigmatism` are `True`. """ def __init__( self, np.ndarray[double, ndim=1] x, np.ndarray[double, ndim=1] y, opd = None, amplitude = None, auto_remove_tilts=False, auto_remove_curvatures=False, auto_remove_astigmatism=False, is_focusing_element=False, put_focal_length=None ): self.x = x.copy() self.y = y.copy() self.X, self.Y = np.meshgrid(x, y) self.R = np.sqrt(self.X**2 + self.Y**2) # stop user from changing the dimensions of the grid self.x.setflags(write=False) self.y.setflags(write=False) self.X.setflags(write=False) self.Y.setflags(write=False) self.is_focusing_element = is_focusing_element self.do_remove_tilts = auto_remove_tilts self.do_remove_curvatures = auto_remove_curvatures self.do_remove_astigmatism = auto_remove_astigmatism self.put_focal_length = put_focal_length if opd is None: self.opd = np.zeros((y.size, x.size), dtype=float) elif callable(opd): self.opd = np.zeros((y.size, x.size), dtype=float) self.has_opd_function = True self._opd_function = opd if isinstance(opd, types.MethodType): self.opd_function_is_method = True else: self.opd_function_is_method = False else: if auto_remove_tilts or auto_remove_curvatures or auto_remove_astigmatism: raise FinesseException(f"Automatic tilt/curvature/astigmatism removal on {self} can only be used when opd is defined by a function, not a fixed array. Use `map.remove_*` instead.") self.has_opd_function = False self.opd = opd.copy() if x.size != self.opd.shape[1] or y.size != self.opd.shape[0]: raise FinesseException("2D opd data and x/y vector shape do not match") # raise exception if map contains any NaNs if np.isnan(self.opd).any(): raise FinesseException("OPD map contains NaN values") if amplitude is None: self.amplitude = np.ones((y.size, x.size), dtype=float) elif callable(amplitude): self.amplitude = np.ones((y.size, x.size), dtype=float) self.has_amplitude_function = True self._amplitude_function = amplitude if isinstance(amplitude, types.MethodType): self.amplitude_function_is_method = True else: self.amplitude_function_is_method = False else: self.has_amplitude_function = False self.amplitude = amplitude.copy() if x.size != self.amplitude.shape[1] or y.size != self.amplitude.shape[0]: raise FinesseException("2D amplitude data and x/y vector shape do not match") # raise exception if map contains any NaNs if np.isnan(self.amplitude).any(): raise FinesseException("Amplitude map contains NaN values") @property def dx(self): return self.x[1] - self.x[0] @property def dy(self): return self.y[1] - self.y[0] def update(self, model=None, weight_spot_size:float=None): """ If the optical path difference or ampltiude data for this map is defined by some callable function you will need to call this update method before accessing its data. If the functions rely on Model parameter values you will also need to provide a Model to evaluate the map with. If automatic tilt or curvature removal is used you will also have to specify the spot size to use as a weighting. Parameters ---------- model : :class:`finesse.Model` Model to use for passing to any map functions. weight_spot_size : float Spot size to use for automatic weighting curvature and tilt removal """ # Update the OPD and amplitudes if they are set by # a callable function, which may or may not be changing, we do # not know currently if self.has_opd_function: if self.opd_function_is_method: self.opd[:, :] = self._opd_function(model) else: self.opd[:, :] = self._opd_function(self, model) # if we have a callable for the opd, we need to remove # any curvature and tilt as required, as a new map might have been # generated. if self.do_remove_tilts: if weight_spot_size is None: raise RuntimeError("Weight spot size must be defined to remove tilts") self.remove_tilts(weight_spot_size) if self.do_remove_curvatures or self.do_remove_astigmatism: self.remove_curvatures(weight_spot_size, "average" if self.do_remove_curvatures else "astigmatic", model=model) if self.has_amplitude_function: if self.amplitude_function_is_method: self.amplitude[:, :] = self._amplitude_function(model) else: self.amplitude[:, :] = self._amplitude_function(self, model) @property def OPD_function(self): """Returns the function that defines the OPD of this map, if there is one. Otherwise it returns `None`.""" return self._opd_function @property def amplitude_function(self): """Returns the function that defines the amplitude of this map, if there is one. Otherwise it returns `None`. """ return self._amplitude_function def rms(self, spot_size : float): """ Computes the spot size weighted RMS of this maps OPD. Parameters ---------- spot_size : float spot size to use for weighting during tilt removal Returns ------- rms : float Root mean square of OPD in metres """ x = np.asarray(self.x) y = np.asarray(self.y) disp = np.asarray(self.opd) return rms(x, y, disp, spot_size) def remove_piston(self, spot_size: float|None): """ Removes any piston term from the optical path difference data. This corresponds to removing a constant term averaged over the spotsize. Note this will only operate on data stored in the ``.opd`` attribute. If the OPD data is determined using a custom function then this method will not work. Instead, you can specify the map to auto-remove any tilts which ensures this method is called after any update to the map is made during the simulation. Parameters ---------- spot_size : float or None spot size to use for weighting. If `None` then a full area average is done. Returns ------- piston : float Removed piston term """ disp = np.asarray(self.opd) if spot_size is not None: if spot_size/self.dx < 10 or spot_size/self.dy < 10: warn("Spot size/dx is low, increase map resolution or spot size") x = np.asarray(self.x) y = np.asarray(self.y) piston = overlap_piston_coefficient(x, y, disp, spot_size) else: piston = np.nanmean(self.opd) disp -= piston return piston def remove_tilts(self, spot_size : float): """ Removes any tilts from the optical path difference data. This corresponds to removing the linear term in the OPD data. Note this will only operate on data stored in the ``.opd`` attribute. If the OPD data is determined using a custom function then this method will not work. Instead, you can specify the map to auto-remove any tilts which ensures this method is called after any update to the map is made during the simulation. Parameters ---------- spot_size : float spot size to use for weighting during tilt removal Returns ------- yaw, pitch : (float, float) Removed linear terms """ if spot_size/self.dx < 10 or spot_size/self.dy < 10: warn("Spot size/dx is low, increase map resolution or spot size") x = np.asarray(self.x) y = np.asarray(self.y) disp = np.asarray(self.opd) yaw, pitch = overlap_tilt_coefficients(x, y, disp, spot_size) # TODO make this more efficient, could do it with 1D arrays X, Y = np.meshgrid(x, y) disp -= (yaw * X + pitch * Y) return yaw, pitch def get_linear_terms(self, spot_size: float): r""" Returns the spot weighted x and y linear terms from the OPD map data. For a surface `z` this is: .. math:: z(x,y) = a x + b y Parameters ---------- spot_size : float Spot size to use for weighting during curvature removal Returns ------- a, b : (float, float) Linear terms in x and y """ if spot_size/self.dx < 10 or spot_size/self.dy < 10: warn("Spot size/dx is low, increase map resolution or spot size") x = np.asarray(self.x) y = np.asarray(self.y) disp = np.asarray(self.opd) a, b = overlap_tilt_coefficients(x, y, disp, spot_size) return a, b def get_quadratic_terms(self, spot_size: float): r""" Returns the spot weighted x and y quadratic terms from the OPD map data. For a surface `z` this is: .. math:: z(x,y) = a x^2 + b y^2 It's common to want to get an equivalent curvature or thin-lens focal length for a map. The quadratic component can be related to the curvature on reflection as .. math:: a = \frac{1}{2 R} or a single pass of a thin lens .. math:: a = -\frac{1}{2 f} Parameters ---------- spot_size : float Spot size to use for weighting during curvature removal Returns ------- a, b : (float, float) Quadratic terms in x and y """ if spot_size/self.dx < 10 or spot_size/self.dy < 10: warn("Spot size/dx is low, increase map resolution or spot size") x = np.asarray(self.x) y = np.asarray(self.y) disp = np.asarray(self.opd) a, b = overlap_curvature_coefficients(x, y, disp, spot_size) return a, b def get_radius_of_curvature_reflection(self, spot_size : float): r""" Assuming this map's OPD array describes a surface, this function will compute the equivalent radius of curvature from it. Uses :meth:`get_quadratic_terms`. Returns ------- Rcx, Rcy Radius of curvature in each direction, in units of meters Notes ----- An ideal parabolic reflecting surface has a surface shape: .. math:: z = \\frac{1}{4 f} r^{2} = a r^{2} where :math:`a` is the quadratic component of the surface, and :math:`f` is the focal length of the reflecting surface. The radius of curvature of a mirror is related to the focal length by :math:`2f=R`, therefore the optical path depth is .. math:: z = \\frac{1}{2 R} r^{2} """ a, b = self.get_quadratic_terms(spot_size) Rc_x = 1/(2*a) Rc_y = 1/(2*b) return Rc_x, Rc_y def get_thin_lens_f(self, spot_size : float, *, average=False): r""" Computes the equivalent thin lens focal length from the optical path depth (OPD) in this map. Uses :meth:`get_quadratic_terms`. Parameters ---------- spot_size : float Circular spot size weighting for curvature removal [m] average : bool, optional Whether to return the average in x and y of the focal length Returns ------- focal_length : tuple|float If `average==False` returns focal_length_x, focal_length_y Focal length in each direction. Otherwise a singular average of the values Notes ----- An ideal thin lens has an optical path length of .. math:: z = -\frac{1}{2 f} r^{2} = a r^{2} where `a` is the quadratic component of the OPD, and `f` is the focal length of the lens. """ a, b = self.get_quadratic_terms(spot_size) if not average: f_a = -1/(2*a) if a != 0 else -np.inf f_b = -1/(2*b) if b != 0 else -np.inf return f_a, f_b else: return -1/(2*(a+b)/2) if a+b != 0 else -np.inf def remove_curvatures(self, spot_size : float, mode:str="average", model=None): """ Removes average weighted quadratic terms from the OPD map data. To get the value of the quadratic term being removed use :meth:`get_quadratic_terms` before calling this. If `put_focal_length` has been set then calling this method will put the removed curvatures into the parameter `put_focal_length`. Parameters ---------- spot_size : float spot size to use for weighting during curvature removal mode : str Different options for curvatures removal are available * "average" : removes average of x and y curvatures * "astigmatic" : removes astigmatic x and y curvatures model : :class:`Model` If `put_focal_length` has been set then it can be used against a specific model object. If `None` it will use the model of the element provided. Returns ------- ax, ay : float Quadratic coefficients in the x and y direction. Same as returned by :meth:`get_quadratic_terms`. """ if self.put_focal_length: fx, fy = self.get_thin_lens_f(spot_size) if model is None: if isinstance(self.put_focal_length, str): raise RuntimeError(f"For the map {repr(self)} to put removed focal length into {repr(self.put_focal_length)} a model must be provided to `update`") else: model = self.put_focal_length._model lens = model.get_element(self.put_focal_length) lens.f.value = (fx+fy)/2 if spot_size/self.dx < 10 or spot_size/self.dy < 10: warn("Spot size/dx is low, increase map resolution or spot size") x = np.asarray(self.x) y = np.asarray(self.y) disp = np.asarray(self.opd) a, b = self.get_quadratic_terms(spot_size) if mode == "average": a = (a+b)/2 b = a elif mode == "astigmatic": a = a b = b else: raise ValueError("Expected `average` or `astigmatic` for removal mode") # TODO make this more efficient, could do it with 1D arrays X, Y = np.meshgrid(x, y) disp -= (a * X**2 + b * Y**2) return a, b def get_z(self, double k, double phase_factor, model=None): r""" Calculates the 2D complex-valued array combining both the amplitude and OPD data stored in this map. This returned value is the 2D data that is integrated over to computed higher order mode scattering matrices. Parameters ---------- k : float Wave number phase_factor : double Optical path difference scaling factor model : :class:`finesse.Model`, optional Model object to use in calculating the map Returns ------- complex_map : ndarray Complex map array Notes ----- Note that if the maps amplitude or OPD data is defined by a custom function that relies on model parameter data you will have to provide a current model to call this function. The phase definition used here for converting OPD to phase is .. math:: \mathrm{amplitude}(x,y) exp(-i k \mathrm{OPD}(x,y)) The ``phase_factor`` sign should be used to convert between different coordinates and types of. For example, `reflections` from port 1 side of a mirror ``phase_factor=-2``. A minus sign is here because the surface normal is pointing back towards the beam, and the 2 is double passing the OPD. """ # Make new array to write complex map data to and return it for user cdef np.ndarray z = np.empty((self.y.size, self.x.size), dtype=complex) # use usual definition of exp(-ikz) then users of this function have to # map from this to whatever coordinate systems they are using np.exp(-1j * k * phase_factor * self.opd, out=z) z *= self.amplitude return z def scatter_matrix(self, q, k0, phase_scale, homs): """Computes a modal scattering matrix for this map. Parameters ---------- q : BeamParam Setting input and output complex beam parameters for overlap calcuations. Options are: - q : all input/output beam parameters are the same - (qx, qy) : astigmatic input/output beam parameters - (q_ix, q_iy, q_ox, q_oy) : general case, each x/y input/out can be different k0 : float Wavenumber of the light to compute scatter for, 2*pi/wavelength phase_scale : float Scaling to apply to the map data to generate the phase information. This sets whether the map is representing a transmission, or reflection. The phase used to compute the scatter is given by: .. math:: exp(-i \textrm{phase}_{\textrm{scale}} k \mathrm{OPD}(x,y)) for example. Some typical values are - phase_scale = 1 : transmission - phase_scale = 2 : reflection homs : array_like A 2D array of (n,m) values for the Hermite-Gaussian modes to compute the scattering for. e.g [[0,0], [1,0], [4,2], ...]. Returns ------- A :class:`KnmMatrix` object with the HOM scattering terms. """ homs = np.asarray(homs) K = map_scattering_coefficients(q, np.max(homs.sum(1)), self.x.copy(), self.y.copy(), self.get_z(k0, phase_scale)) knm = scattering_coefficients_to_KnmMatrix(homs, K) return knm def map_scattering_coefficients( q : BeamParam, int max_order, double[::1] x, double[::1] y, complex[:,::1] Z, bint reverse_gouy=False ): r""" Calculates the mode scattering coefficients up to some maximum mode order for a uniformly spaced 2-dimensional map. Essentially it is computing multiple overlap integrals between each input and output mode for some complex valued abberation map, Z. .. math:: K_{abnm} = \iint^{\infty}_{-\infty} U_{nm}(x,y,q_{ix},q_{iy}) Z(x,y) U^*_{ab}(x,y,q_{ox},q_{oy}) \, dy \, dx - ab are the output mode indicies - nm are the input mode indicies Parameters ---------- q : BeamParam Setting input and output complex beam parameters for overlap calcuations. Options are: - q : all input/output beam parameters are the same - (qx, qy) : astigmatic input/output beam parameters - (q_ix, q_iy, q_ox, q_oy) : general case, each x/y input/out can be different max_order : int Maximum order of mode scattering to compute up to. x, y : array[double] 1-dimensional arrays specifying the x and y uniformally sampled axis for the map Z : array[complex] 2-dimensional array describing the complex abberation being applied to some input beam. The shape of this array should be Z[Ny, Nx], where Ny and Nx are the number of y and x samples respectively. reverse_gouy : bool, optional If true, the gouy phase terms are removed from the scattering coefficients Returns ------- K : array[complex] Coupling coefficients array with indexing (inx, outx, iny, outy) upto max_order input. Examples -------- Same input and output beam parameter but computing the scattering from some arbitrary map:: import finesse from finesse.knm.maps import map_scattering_coefficients import numpy as np maxtem = 10 q = finesse.BeamParam(w0=1e-3, z=0) x = np.linspace(-4*q.w, 4*q.w, 201) y = np.linspace(-4*q.w, 4*q.w, 200) X, Y = np.meshgrid(x, y) # Mix of liner and quadratic phase terms to compute scattering # coefficients for Z = np.exp( 1j * 1e-6 * 2 * np.pi * (0.2*X + 0.1*Y + 1000*X**2 - 1000*Y**2) / 1064e-9 ) K = map_scattering_coefficients(q, maxtem, x, y, Z) print("HG 10 -> 30", K[1,3,0,0]) print("HG 00 -> 00", K[0,0,0,0]) print("HG 00 -> 21", K[0,2,0,1]) Different input and output beam parameters. Here modelling mode mismatch between different beam parameters:: qx1 = finesse.BeamParam(w0=1e-3, z=0) qy1 = qx1 qx2 = finesse.BeamParam(w0=1.2e-3, z=0) qy2 = qx2 Z = np.ones_like(Z) K = map_scattering_coefficients((qx1,qy1,qx2,qy2), maxtem, x, y, Z) print("HG 00 -> 20", K[0,2,0,0]) print("HG 00 -> 00", K[0,0,0,0]) """ if not (Z.shape[0] == y.size and Z.shape[1] == x.size): # map must be of shape (Ny, Nx) raise RuntimeError("Z must be of size [y.size, x.size]") try: # Have all four qs qix, qiy, qox, qoy = q except TypeError: # no len option? only q given qix = qiy = qox = qoy = q except ValueError: # wrong size, try: # only provide (qx, qy) qix, qiy = q qox, qoy = q except ValueError: # single q value but in list qix = qiy = qox = qoy = q[0] modes = np.array(tuple( ( (n, m) for n in range(max_order + 1) for m in range(max_order + 1) if n + m <= max_order ) ), dtype=np.int32) dx = x[1] - x[0] dy = y[1] - y[0] w = min((qix.w, qiy.w, qox.w, qoy.w)) if w/dx < 10 or w/dy < 10: warn("Spot size vs map resolution is low, increase map resolution or spot size") # astigmatic = (qix.q != qiy.q) or (qox.q != qoy.q) mismatched = (qix.q != qox.q) or (qiy.q != qoy.q) iHGs = HGModes((qix, qiy), modes, reverse_gouy=reverse_gouy) Uni, Umi = iHGs.compute_1d_modes(x, y) if mismatched: oHGs = HGModes((qox, qoy), modes, reverse_gouy=reverse_gouy) Uno, Umo = oHGs.compute_1d_modes(x, y) Nm = max_order + 1 Umm_ = np.zeros((Nm, Nm, y.size), dtype=np.complex128) Unn_ = np.zeros((Nm, Nm, x.size), dtype=np.complex128) tmp = np.zeros((Nm, Nm, y.size), dtype=np.complex128) K = np.zeros((Nm, Nm, Nm, Nm), dtype=np.complex128) if mismatched: outer_conj_product_2(Umi, Umo, Umm_) outer_conj_product_2(Uni, Uno, Unn_) else: outer_conj_product(Umi, Umm_) outer_conj_product(Uni, Unn_) compute_map_scattering_coeffs_riemann_optimised(dx * dy, Z, Unn_, Umm_, tmp, K) return np.transpose(K, (1, 2, 0, 3)) # change to (inx, outx, iny, outy) def scattering_coefficients_to_KnmMatrix(modes, K): """ Converts a 4-D scattering coefficient array into a 2D KnmMatrix object to be used in simulations. Parameters ---------- modes : tuple, array Array of 2D modes indicies (n,m) to specify the order in which they appear in the returned matrix. K : array 4D array of coefficients indexed with [in_x, out_x, in_y, out_y] Returns ------- Knm : KnmMatrix A KnmMatrix class representing a 2D scattering matrix for the requested modes. Examples -------- Compute KNM matrix from a previously calculated 4-D scattering coefficient array:: import finesse from finesse.knm.maps import ( map_scattering_coefficients, scattering_coefficients_to_KnmMatrix ) import numpy as np # compute all scatterings up to maximum TEM order maxtem = 3 q = finesse.BeamParam(w0=1e-3, z=0) x = np.linspace(-4*q.w, 4*q.w, 201) y = np.linspace(-4*q.w, 4*q.w, 200) X, Y = np.meshgrid(x, y) # Mix of liner and quadratic phase terms to compute scattering # coefficients for Z = np.exp(1j * 1e-6 * 2 * np.pi * X / 1064e-9) K = map_scattering_coefficients(q, maxtem, x, y, Z) # Generate specific ordering of coefficients into a 2D matrix that can # be used in simulations for propagating and array of modes. modes = ( (0,0), (1,0), (2,1), ) Knm = scattering_coefficients_to_KnmMatrix(modes, K) # Propagate some mode amplitudes Knm.data @ [1, 0.1, 0] >>> array([9.99995641e-01+2.95261180e-04j, 9.99986923e-02+2.95261180e-03j, 4.23516474e-22+2.16840434e-20j]) """ from finesse.knm.matrix import KnmMatrix Nm = len(modes) assert K.ndim == 4 _modes = np.array(modes, dtype=np.int32) Knm = np.zeros((Nm, Nm), dtype=complex) for i, (n, m) in enumerate(modes): for o, (a, b) in enumerate(modes): Knm[o, i] = K[n, a, m, b] return KnmMatrix.from_buffer(Knm, _modes) @cython.boundscheck(False) @cython.wraparound(False) @cython.initializedcheck(False) cdef void c_scattering_coefficients_to_KnmMatrix( int[:,::1] modes, Py_ssize_t Nm, complex_t *K, DenseZMatrix *out ) noexcept nogil: """ Converts a 4-D scattering coefficient array into a 2D KnmMatrix object to be used in simulations. Parameters ---------- modes : tuple, array Array of 2D modes indicies (n,m) to specify the order in which they appear in the returned matrix. Nm : int Number of modes computed in K K : *complex_t Flatten 4D array of coefficients [Nm, Nm, Nm, Nm]. Assumes same indexation as computed by compute_map_scattering_coeffs_riemann_optimised. Which of writing is [m, a, n, b] where (a,b) are the output and (n,m) are the input out : *DenseZMatrix A view to a KnmMatrix class' DenseZMatrix representing a 2D scattering matrix for the requested modes """ cdef: Py_ssize_t i, o, n, m, a, b # We have a flat representation so need to multiply # by strides set by Nm for i in range(out.size2): n = modes[i, 0] * Nm*Nm m = modes[i, 1] * Nm*Nm*Nm for o in range(out.size1): a = modes[o, 0] * Nm b = modes[o, 1] out.ptr[o * out.size2 + i] = K[m + n + a + b] cdef update_map_data_in_workspace(knm_map_workspace *ws) : """For a given map workspace this method will recalculate the array used for integrations. This will call any map callable methods set. """ cdef Map _map = <Map>ws.map_obj # use average spot size cdef double spot_size = 0.5*(bp_beamsize(&ws.q_from.qx) + bp_beamsize(&ws.q_from.qy)) _map.update(<object>ws.model, spot_size) cdef complex[:, ::1] z = _map.get_z( ws.k, ws.phase_factor, <object>ws.model ) if z.shape[0] != ws.Ny or z.shape[1] != ws.Nx: raise RuntimeError(f"Updated map has wrong dimensions. Must be the same dimension as original map ({ws.Ny}, {ws.Nx})") # copy map data across to memory we have already allocated memcpy(ws.z, &z[0,0], sizeof(complex_t)*ws.Nx*ws.Ny) ws.new_map_data = True # set flag to recompute maps cdef set_knm_map_workspace( object model, knm_map_workspace *ws, Map _map, double k, double phase_factor ) : """ Set the knm map workspace to link to the Map object. Must call free_knm_map_workspace once finished to decrease reference count to _map. """ assert(ws != NULL) assert(_map is not None) # Need to set as writeable to make a memoryview of array # not really a problem as this is just to stop the user # messing around with the values after they have been set _map.x.setflags(write=True) _map.y.setflags(write=True) cdef: double[::1] _x = _map.x double[::1] _y = _map.y _map.x.setflags(write=False) _map.y.setflags(write=False) ws.map_obj = <PyObject*>_map ws.is_focusing_element = _map.is_focusing_element ws.new_map_data = True Py_XINCREF(ws.map_obj) ws.phase_factor = phase_factor ws.k = k ws.x = &_x[0] ws.y = &_y[0] ws.dx = abs(ws.x[1] - ws.x[0]) ws.dy = abs(ws.y[1] - ws.y[0]) ws.Nx = _map.x.size ws.Ny = _map.y.size # need reference to model so later we can update changing maps ws.model = <PyObject*>model # create empty array for storing data in cdef np.ndarray z = np.zeros((ws.Ny, ws.Nx), dtype=complex) cdef complex_t[:, ::1] _z = z ws.z_obj = <PyObject*>z # keep reference so array memory doesn't get # freed whilst we're using it Py_XINCREF(ws.z_obj) ws.z = &_z[0,0] ws.map_is_changing = _map.has_amplitude_function or _map.has_opd_function cdef init_knm_map_workspace( knm_map_workspace *ws, int Nm, NodeBeamParam *q_from, NodeBeamParam *q_to, bint flip_odd_x_output_modes ) : """Allocates memory for map calculations. Must call free_knm_map_workspace once finished to release memory. Nm : int Number of 1D modes to compute, essentially maxtem+1 """ if Nm <= 0: raise RuntimeError("Nm <= 0") ws.Nm = Nm if ws.Un != NULL: raise MemoryError() ws.Un = <complex_t*> calloc(Nm * ws.Nx, sizeof(complex_t)) if not ws.Un: raise MemoryError() if ws.Um != NULL: raise MemoryError() ws.Um = <complex_t*> calloc(Nm * ws.Ny, sizeof(complex_t)) if not ws.Um: raise MemoryError() if ws.Un_ != NULL: raise MemoryError() ws.Un_ = <complex_t*> calloc(Nm * ws.Nx, sizeof(complex_t)) if not ws.Un_: raise MemoryError() if ws.Um_ != NULL: raise MemoryError() ws.Um_ = <complex_t*> calloc(Nm * ws.Ny, sizeof(complex_t)) if not ws.Um_: raise MemoryError() if ws.Unn_ != NULL: raise MemoryError() ws.Unn_ = <complex_t*> calloc(Nm * Nm * ws.Nx, sizeof(complex_t)) if not ws.Unn_: raise MemoryError() if ws.Umm_ != NULL: raise MemoryError() ws.Umm_ = <complex_t*> calloc(Nm * Nm * ws.Ny, sizeof(complex_t)) if not ws.Umm_: raise MemoryError() if ws.tmp != NULL: raise MemoryError() ws.tmp = <complex_t*> calloc(Nm * Nm * ws.Ny, sizeof(complex_t)) if not ws.tmp: raise MemoryError() if ws.K != NULL: raise MemoryError() ws.K = <complex_t*> calloc(Nm * Nm * Nm * Nm, sizeof(complex_t)) if not ws.K: raise MemoryError() if ws.uiws != NULL: raise MemoryError() ws.uiws = <unm_workspace*> calloc(1, sizeof(unm_workspace)) if not ws.uiws: raise MemoryError() if ws.uows != NULL: raise MemoryError() ws.uows = <unm_workspace*> calloc(1, sizeof(unm_workspace)) if not ws.uows: raise MemoryError() if ws.unm_i_factor_ws != NULL: raise MemoryError() ws.unm_i_factor_ws = <unm_factor_store*> calloc(1, sizeof(unm_factor_store)) if not ws.unm_i_factor_ws: raise MemoryError() if ws.unm_o_factor_ws != NULL: raise MemoryError() ws.unm_o_factor_ws = <unm_factor_store*> calloc(1, sizeof(unm_factor_store)) if not ws.unm_o_factor_ws: raise MemoryError() ws.q_to = q_to ws.q_from = q_from unm_ws_recache_from_bp(ws.uiws, &ws.q_from.qx, &ws.q_from.qy) unm_ws_recache_from_bp(ws.uows, &ws.q_to.qx, &ws.q_to.qy) unm_factor_store_init(ws.unm_i_factor_ws, ws.uiws, Nm, Nm, True, False) # Only flip output x modes unm_factor_store_init(ws.unm_o_factor_ws, ws.uows, Nm, Nm, True, flip_odd_x_output_modes) # as all q values have been set we can compute # the map array and load it into the workspace update_map_data_in_workspace(ws) cdef free_knm_map_workspace(knm_map_workspace *ws) : Py_XDECREF(ws.map_obj) Py_XDECREF(ws.z_obj) if ws.Un: free(ws.Un) if ws.Un_: free(ws.Un_) if ws.Um: free(ws.Um) if ws.Um_: free(ws.Um_) if ws.Unn_: free(ws.Unn_) if ws.Umm_: free(ws.Umm_) if ws.tmp: free(ws.tmp) if ws.K: free(ws.K) if ws.unm_o_factor_ws: free(ws.unm_o_factor_ws) if ws.unm_i_factor_ws: free(ws.unm_i_factor_ws) if ws.uows: free(ws.uows) if ws.uiws: free(ws.uiws)