Backends¶

SAX Backends

Introduction¶

SAX allows to easily interchange the backend of a circuit. A SAX backend consists of two static analysis steps and an evaluation step:

:::{eval-rst} .. autofunction:: sax.backends.analyze_instances :::

:::{eval-rst} .. autofunction:: sax.backends.analyze_circuit :::

:::{eval-rst} .. autofunction:: sax.backends.evaluate_circuit :::

The analyze_instances step analyzes the 'shape' of the instances by running each model with default parameters. By shape we mean which port-combinations are present in the sparse S-matrix.

QUESTION: can we do this more efficiently in a functional way?

After analyzing the instances, it is assumed that the shape of the instance won't change any more. It is therefore important that you write your model functions in such a way that the port combinations present in your s-matrix never changes!

NOTE: we used to do this analysis step in analyze_circuit by just looking at the connections present. However, this inherently assumed dense connectivity within the model. This is pretty inefficient for large sparse models with many ports. Ideally, however, we should be able to analyze the shape of our instances without running their models with default parameters...

The analyze_circuit step should statically analyze the connections and ports and should return an analyzed object. This object contains all the static objects that are needed for circuit computation but won't be recalculated when any parameters of the circuit change. See KLU backend for a non-trivial implementation of the circuit analyzation.

The evaluate_circuit step evaluates the circuit for given SType instances, given whatever analysis object was returned from the analyze_circuit step and the instance STypes.

Example¶

Filipsson-Gunnar Backend¶

:::{note} Since The KLU Backend (see below) is Superior to the Filipsson-Gunnar backend, SAX will default (since v0.10.0) to the KLU backend if klujax is installed.name1 :::

The Filipsson-Gunnar backend is based on the following paper:

Filipsson, Gunnar. "A new general computer algorithm for S-matrix calculation of interconnected multiports." 11th European Microwave Conference. IEEE, 1981.

:::{eval-rst} .. autofunction:: sax.backends.analyze_circuit_fg :::

:::{eval-rst} .. autofunction:: sax.backends.evaluate_circuit_fg :::

Algorithm Walkthrough¶

:::{note} This algorithm gets pretty slow for large circuits. Since SAX v0.10.0 we will default to the superior KLU backend as the KLU backend is now also jittable. :::

Let's walk through all the steps of this algorithm. We'll do this for a simple MZI circuit, given by two directional couplers characterised by dc_sdense with a phase shifting waveguide in between wg_sdict:

as a first step, we construct the reversed_ports, it's actually easier to work with reversed_ports (we chose the opposite convention in the netlist definition to adhere to the GDSFactory netlist convention):

The first real step of the algorithm is to create the 'block diagonal sdict`:

we can optionally filter out zeros from the resulting block_diag representation. Just note that this will make the resuling function unjittable (the resulting 'shape' (i.e. keys) of the dictionary would depend on the data itself, which is not allowed in JAX jit). We're doing it here to avoid printing zeros but internally this is not done by default.

next, we sort the connections such that similar components are grouped together:

Now we iterate over the sorted connections and connect components as they come in. Connected components take over the name of the first component in the connection, but we keep a set of components belonging to that key in all_connected_instances.

This is how this all_connected_instances dictionary looks initially.

Normally we would loop over every connection in sorted_connections now, but let's just go through it once at first:

k and l are the S-matrix indices we're trying to connect. Note that in our sparse SDict notation these S-matrix indices are in fact equivalent with the port names ('dc1,out1', 'dc2,in1')!

first we split the connection string into an instance name and a port name (we don't use the port name yet):

We then obtain the new set of connected instances.

We then iterate over each of the components in this set and make sure each of the component names in that set maps to that set (yes, I know... confusing). We do this to be able to keep track with which components each of the components in the circuit is currently already connected to.

now we need to obtain all the ports of the currently connected instances.

Now the Gunnar Algorithm is used. Given a (block-diagonal) 'S-matrix' block_diag and a 'connection matrix' current_ports we can interconnect port k and l as follows:

Note: some creative freedom is used here. In SAX, the matrices we're talking about are in fact represented by a sparse dictionary (an SDict), i.e. similar to a COO sparse matrix for which the indices are the port names.

Just as before, we're filtering the zeros from the sparse representation (remember, internally this is not done by default).

This is the resulting block-diagonal matrix after interconnecting two ports (i.e. basically saying that those two ports are the same port). Because these ports are now connected we should actually remove them from the S-matrix representation (they are integrated into the S-parameters of the other connections):

Note that this deletion of values does NOT make this operation un-jittable. The deletion depends on the ports of the dictionary (i.e. on the dictionary 'shape'), not on the values.

We now basically have to do those steps again for all other connections:

This is the final MZI matrix we're getting:

All that's left is to rename these internal ports of the format {instance},{port} into output ports of the resulting circuit:

And that's it. We evaluated the SDict of the full circuit.

Algorithm Improvements¶

The Filipsson-Gunar algorithm is

• pretty fast for small circuits 🙂
• jittable 🙂
• differentiable 🙂
• GPU-compatible 🙂

This algorithm is however:

• really slow for large circuits 😥
• pretty slow to jit the resulting circuit function 😥
• pretty slow to differentiate the resulting circuit function 😥

There are probably still plenty of improvements possible for this algorithm:

• ¿ Network analysis (ft. NetworkX ?) to obtain which ports of the block diagonal representation are relevant to obtain the output connection ?
• ¿ Smarter ordering of connections to always have the minimum amount of ports in the intermediate block-diagonal representation ?
• ¿ Using jax.lax.scan in stead of python native for-loops in _interconnect_ports ?
• ¿ ... ?

Bottom line is... Do you know how to improve this algorithm or how to implement the above suggestions? Please open a Merge Request!

KLU Backend¶

The KLU backend is using klujax, which uses the SuiteSparse C++ libraries for sparse matrix evaluations to evaluate the circuit insanely fast on a CPU. The specific algorith being used in question is the KLU algorithm:

Ekanathan Palamadai Natariajan. "KLU - A high performance sparse linear solver for circuit simulation problems."

:::{eval-rst} .. autofunction:: sax.backends.analyze_circuit_klu :::

:::{eval-rst} .. autofunction:: sax.backends.evaluate_circuit_klu :::

Theoretical Background¶

The core of the KLU algorithm is supported by klujax, which internally uses the Suitesparse libraries to solve the sparse system Ax = b, in which A is a sparse matrix.

Now it only comes down to shoehorn our circuit evaluation into a sparse linear system of equations $Ax=b$ where we need to solve for $x$ using klujax. Consider the block diagonal matrix $S_{bd}$ of all components in the circuit acting on the fields $x_{in}$ at each of the individual ports of each of the component integrated in $S^{bd}$. The output fields $x^{out}$ at each of those ports is then given by:

$$ x^{out} = S_{bd} x^{in} $$

However, $S_{bd}$ is not the S-matrix of the circuit as it does not encode any connectivity between the components. Connecting two component ports basically comes down to enforcing equality between the output fields at one port of a component with the input fields at another port of another (or maybe even the same) component. This equality can be enforced by creating an internal connection matrix, connecting all internal ports of the circuit:

$$ x^{in} = C_{int} x^{out} $$

We can thus write the following combined equation:

$$ x^{in} = C_{int} S_{bd} x^{in} $$

But this is not the complete story... Some component ports will not be interconnected with other ports: they will become the new external ports (or output ports) of the combined circuit. We can include those external ports into the above equation as follows:

$$ \begin{pmatrix} x^{in} \\ x^{out}_{ext} \end{pmatrix} = \begin{pmatrix} C_{int} & C_{ext} \\ C_{ext}^T & 0 \end{pmatrix} \begin{pmatrix} S_{bd} x^{in} \\ x_{ext}^{in} \end{pmatrix} $$

Note that $C_{ext}$ is obviously not a square matrix. Eliminating $x^{in}$ from the equation above finally yields:

$$ x^{out}_{ext} = C^T_{ext} S_{bd} (I - C_{int}S_{bd})^{-1} C_{ext}x_{ext}^{in} $$

We basically found a representation of the circuit S-matrix:

$$ S = C^T_{ext} S_{bd} (I - C_{int}S_{bd})^{-1} C_{ext} $$

Obviously, we won't want to calculate the inverse $(I - C_{int}S_{bd})^{-1}$, which is the inverse of a very sparse matrix (a connection matrix only has a single 1 per line), which very often is not even sparse itself. In stead we'll use the solve_klu function:

$$ S = C^T_{ext} S_{bd} \texttt{solve}\_\texttt{klu}\left((I - C_{int}S_{bd}), C_{ext}\right) $$

Moreover, $C_{ext}^TS_{bd}$ is also a sparse matrix, therefore we'll also need a mul_coo routine:

$$ S = C^T_{ext} \texttt{mul}\_\texttt{coo}\left(S_{bd},~~\texttt{solve}\_\texttt{klu}\left((I - C_{int}S_{bd}),~C_{ext}\right)\right) $$

Sparse Helper Functions¶

:::{eval-rst} .. autofunction:: klujax.solve :::

:::{eval-rst} .. autofunction:: klujax.coo_mul_vec :::

klujax.solve solves the sparse system of equations Ax=b for x. Where A is represented by in COO-format as (Ai, Aj, Ax).

Example

This result is indeed correct:

However, to use this function effectively, we probably need an extra dimension for Ax. Indeed, we would like to solve this equation for multiple wavelengths (or more general, for multiple circuit configurations) at once. For this we can use jax.vmap to expose klujax.solve to more dimensions for Ax:

Let's now redefine Ax and see what it gives:

This result is indeed correct:

Additionally, we need a way to multiply a sparse COO-matrix with a dense vector. This can be done with klujax.coo_mul_vec:

However, it's useful to allow a batch dimension, this time both in Ax and in b:

Let's confirm this does the right thing:

Example¶

the KLU backend yields SDense results by default:

An SDense is returned for perfomance reasons. By returning an SDense by default we prevent any internal SDict -> SDense conversions in deeply hierarchical circuits. It's however very easy to convert SDense to SDict as a final step. To do this, wrap the result (or the function generating the result) with sdict:

Algorithm Walkthrough¶

Let's first enforce $C^T = C$:

We'll also need the reversed ports:

An the port indices

Let's now create the COO-representation of our block diagonal S-matrix $S_{bd}$:

note that we also kept track of the batch_shape, i.e. the number of independent simulations (usually number of wavelengths). In the example being used here we don't have a batch dimension (all elements of the SDict are 0D):

We'll also keep track of the number of columns

And we'll need to solve the circuit for each output port, i.e. we need to solve n_rhs number of equations:

We can represent the internal connection matrix $C_{int}$ as a mapping between port indices:

Therefore, the COO-representation of this connection matrix can be obtained as follows (note that an array of values Cx is not necessary, all non-zero elements in a connection matrix are 1)

We can represent the external connection matrix $C_{ext}$ as a map between internal port indices and external port indices:

Just as for the internal matrix we can represent this external connection matrix in COO-format:

However, we actually need it as a dense representation:

help needed: can we find a way later on to keep this sparse?

We'll now calculate the row index CSi of $C_{int}S_{bd}$ in COO-format:

CSi: possible jittable alternative? how do we remove the zeros?

The column index CSj of $C_{int}S_{bd}$ can more easily be obtained:

Finally, the values CSx of $C_{int}S_{bd}$ can be obtained as follows:

Now we calculate $I - C_{int}S_{bd}$ in an uncoalesced way (we might have duplicate indices on the diagonal):

uncoalesced: having duplicate index combinations (i, j) in the representation possibly with different corresponding values. This is usually not a problem as in linear operations these values will end up to be summed, usually the behavior you want:

The batch shape dimension can generally speaking be anything (in the example here 0D). We need to do the necessary reshapings to make the batch shape 1D:

We're finally ready to do the most important part of the calculation, which we conveniently leave to klujax and SuiteSparse:

one more sparse multiplication:

And one more $C_{ext}$ multiplication which we do by clever indexing:

That's it! We found the S-matrix of the circuit. We just need to reshape the batch dimension back into the matrix:

Oh and to complete the SDense representation we need to specify the port map as well:

Algorithm Improvements¶

This algorithm is

• very fast for large circuits 🙂
• jittable 🙂
• differentiable 🙂

This algorithm is however:

• not GPU-compatible 🙂

There are probably still plenty of improvements possible for this algorithm:

• ¿ make it GPU compatible (requires making suitesparse GPU compatible... probably not gonna happen)?

:::{note} Since The KLU Backend is Superior to the Filipsson-Gunnar backend, SAX will default (since v0.10.0) to the KLU backend if klujax is installed.name1 :::

Additive Backend¶

:::{eval-rst} .. autofunction:: sax.backends.analyze_circuit_additive :::

:::{eval-rst} .. autofunction:: sax.backends.evaluate_circuit_additive :::

Sometimes we would like to calculate circuit path lengths or time delays within a circuit. We could obviously simulate these things with a time domain simulator, but in many cases a simple additive backend (as opposed to the default multiplicative backend) can suffice.

Example¶

:::{note} in stead of S-parameters the stypes need to contain additive parameters, such as length or time delay. :::

:::{note} it is recommended to not use an SDense representation for the additive backend. Very often an SDense representation will introduce zeros which will be interpreted as an existing connection with zero length. Conversely, in a sparse representation like SDict or SCoo, non-existing elements will be just that: they will not be present in the internal graph. :::

We made a difference here between edges of 'S'-type (connections through the S-matrix) and edges of 'C'-type (connections through the connection matrix). Connections of 'C'-type obviously always have length zero as they signify per definition the equality of two ports.

We can create a NetworkX graph from these edges:

We can now get a list of all possible paths in the network. Note that these paths must alternate between an S-edge and a C-edge:

And the path lengths of those paths can be calculated as follows:

This is all brought together in the additive KLU backend: