import numpy as np
from ...solutions import BaseSolution
from .base import Action
[docs]class DCFieldsSolutions(BaseSolution):
"""Contains the result of a :class:`DCFields` action.
Attributes
----------
homs : list
The Hermite-Gaussian higher order mode indices at each node, of each frequency.
nodes : tuple[str]
Name of nodes the fields were extracted at
frequencies : array_like[dtype=float]
Values of the optical frequencies at each node
fields : array_like
Field data array with shape [nodes, frequencies, HOMs.shape[0]]
Notes
-----
The `fields` attribute will contain all the required information with its ordering
defined by the `nodes`, `frequencies`, and `homs` array. For easier indexing you can
access the solution directly with node names.
Select all the frequencies and modes at l1.p1.i:
>>> sol['l1.p1.i']
Select all the first modes at all frequencies and modes at l1.p1.i:
>>> sol['l1.p1.i', :, 0]
Select all the frequencies and modes at two nodes:
>>> sol[('l1.p1.i', 'l1.p1.o')]
Select all the HG02 modes at every node and frequency
>>> sol[:, :, sol.homs.index([0, 2]))]
Select all the HG02 modes at every output node and every frequency
>>> idx = np.array([node.split(".")[-1] == "o" for node in sol.nodes])
>>> sol[idx, :, sol.homs.index([0, 2])]
The frequency and mode selection must be a slice or an array of integer or boolean
values.
"""
homs = None
nodes = None
frequencies = None
fields = None
def __getitem__(self, key):
if isinstance(key, (str, int, slice, np.ndarray)):
# specifying a single node or a slice of multiple nodes
key = [key]
elif isinstance(key, tuple):
# this is either a tuple specifying multiple nodes or
# specifying multiple dimensions
assert len(key) > 1
if len(key) in [2, 3]:
if isinstance(key[1], str):
# specifying multiple nodes
key = [key]
else:
# specifying multiple dimensions
key = list(key)
else:
# specifying multiple nodes
key = [list(key)]
else:
raise TypeError(
f"Unrecognized key type {type(key)}."
"Multiple nodes should be specified by tuples, slices, or ndarrays."
)
if not isinstance(key[0], (slice, int, np.ndarray)):
nodes = np.atleast_1d(key[0])
for i in nodes:
if i not in self.nodes:
raise KeyError(f"`{i}` is not in the solution nodes")
key[0] = tuple(self.nodes.index(i) for i in nodes)
return self.fields[tuple(key)]
[docs]class DCFields(Action):
"""An action that saves the DC (carrier) fields at all nodes, frequencies, and
higher order modes for the current state of the simulation.
Parameters
----------
name : str, optional
Name of the solution generated by this action
"""
def __init__(self, *, name="dcfields"):
super().__init__(name)
def _requests(self, model, memo, first=True):
pass
def _do(self, state):
sol = DCFieldsSolutions(self.name)
sol.homs = state.model.homs.tolist()
sol.nodes = tuple(state.sim.carrier.nodes.keys())
sol.frequencies = np.array(
[f.f for f in state.sim.carrier.optical_frequencies.frequencies]
)
sol.fields = np.reshape(
state.sim.carrier.M().rhs_view[0, :].copy(),
(len(sol.nodes), len(sol.frequencies), len(sol.homs)),
order="C",
)
# rescale into sqrt(W) units
sol.fields /= np.sqrt(2)
return sol