HOM performance improvements
Link to discussions on mattermost: https://chat.ligo.org/ligo/pl/3nobkuswdfdft8mqfphnrboyhr
Following on from the general / plane-wave performance enhancements completed in !17 (merged), several speed-up improvements should now be made to HOM-related code. The task list below indicates some likely areas to look at.
-
Store beam parameters in a C array of structs in the BaseSimulation
class. -
Change the way ABCD matrices are stored in components. -
Implement KnmWorkspace
extension class infinesse.knm
to avoid repeated calculations in BH scatter matrix code. -
Cache u_nm
distribution and Gouy phases in camera extension - completed in 38a41758, e0b5bc04 and 205dde1c. -
Cythonise all _compute_knm_matrices
methods (and any dependent methods) so that aprange
loop over all knm matrices can be tried out. -
Fast version of beam tracing on frozen model.
Beam parameters
Currently the recipe for retrieving beam parameters during a simulation is:
node_q_map = sim.model.last_trace
qx, qy, _ = node_q_map[node]
This requires a fair amount of Python interaction due to dictionary access and attribute retrieval. To eliminate this, beam parameters should be stored in some way in the SimulationTraceSolution
or BaseSimulation
itself. I recently started implementing this in the NodeInfoEntry
struct of simulations/base.pyx
but it may be better to store in a separate struct as suggested by @daniel-brown via something like:
cdef struct NodeBeamParam:
complex_t qx
complex_t qy
bint is_changing
These could then be stored as a C array (with indices initialised using node ids via BaseSimulation.node_id
similarly to how the _c_node_info
array works) in BaseSimulation
, e.g:
cdef class BaseSimulation:
cdef:
# ...
NodeBeamParam* trace
# ...
or as some attribute of the SimulationTraceSolution
class. At the moment I prefer the former as it would cut out all Python interaction for retrieving beam params - which would be required if we want a parallel for loop over scattering matrix computations for example. Alternatively, the SimulationTraceSolution
class could also be Cythonised but that might be overkill.
The construction of SimulationTraceSolution
objects could be cut out completely in Model.beam_trace
when the model is built (i.e. in a simulation state) and instead the beam parameters at nodes can be updated directly via something like:
# in the Tracer class in __simtrace_step method
if self.model.is_built:
self.model.carrier_simulation.update_beam_parameter(node_id, qx, qy)
else:
# just assign to the dictionary as we do currently
Storing ABCD matrices
ABCD matrices are computed via the relevant Connector.ABCD
method where this method will cache a numeric ABCD on first calculation (storing it in the Connector._abcd_matrices
dict). If a dependent parameter changes then the Connector._recompute_abcd
dict will switch the couplings to True to indicate that the next call to the ABCD method should re-compute the matrix rather than retrieving the current cached one.
This behaviour should now be changed. My current thinking for this is that symbolic ABCD matrices should be created for each optical coupling of the component on construction of the component. These will then be stored (in Connector._abcd_matrices
dict?). If the component is a Surface
(i.e. the ABCD matrices depend upon refractive indices of adjacent spaces) then the symbolic ABCD matrices will get re-constructed when adding it to a model (or removing it / connected spaces).
Upon the start of a simulation, each components' ABCD matrices should be evaluated (cast the array to floats) and these would then be stored as memory-views in the relevant connector workspaces. For example, the mirror would have something like:
cdef class MirrorWorkspace:
cdef:
# ...
double[:, ::1] abcd_p1p1_x # reflection at port p1 in tangential plane
double[:, ::1] abcd_p1p1_y # reflection at port p2 in sagittal plane
double[:, ::1] abcd p1p2_x # transmission from p1 -> p2 in tangential plane
# etc.
Something that still needs to be figured out is how to efficiently update these ABCD matrices when a dependent parameter changes (e.g. when an xaxis scan is over a mirror RoC).
KnmWorkspace
Connector workspaces should also store KnmWorkspace
extension (cdef
) classes for each coupling. These store common parameters required for each element of the scattering matrices for each coupling, and should take care of determining whether a Knm matrix needs to be re-computed.
This feature has mostly been completed but not yet pushed.
u_nm
distribution caching
Camera See 38a41758, e0b5bc04 and 205dde1c for details on the completed implementation of this performance enhancement.
_compute_knm_matrices
Cythonising As the _compute_knm_matrices
methods are all Python methods of the relevant Connector
classes, currently a Python for loop over all these is performed in BaseSimulation
:
cpdef compute_knm(self, bint force=False):
for fn in self._compute_knm_matrices:
fn(self, force)
A potentially massive performance improvement (especially for simulations with changing RoCs / lengths / x,y betas) would be to Cythonise all of these methods so that we can do something like:
cdef compute_knm(self):
# N is number of components which scatter modes
for i in prange(N, nogil=True):
# some code to get relevant workspace or something
ws.compute_scatter_matrices()
This kind of feature will require direct C access to beam parameters, ABCD matrices and refractive indices.
Fast beam tracing
It may turn out that the beam tracing itself could become a bottleneck in some simulations - most likely those which don't require matrix solving (i.e. for models containing only bp
, cp
, gouy
detectors). If this is the case, then a faster version of the beam tracing should be implemented along-side the current algorithm. This new version would be Cythonised and rely on accessing connector workspaces, it would also use the frozen model to explicitly pre-determine only the paths which will have changing beam parameters - this could be determined via modification to the current Model.beam_trace
or a separate function (either way, this pre-determination of "changing paths" would need to happen at the start of the simulation).
Note that this re-factoring could be worth doing anyway (even if beam tracing is not a bottleneck) as it could serve as a clean way to separate beam tracing in a simulation from beam tracing performed outside of a simulation.