DriveParameters -> Model, which should naturally include the Hamiltonian

docs/floquet.md

@@ -14,9 +14,9 @@ The **floquet** Python API consists largely of the **floquet_analysis** function
     options:
         show_source: false
 
-## Drive parameters
+## Model
 
-::: floquet.drive_parameters
+::: floquet.model
     options:
         show_source: false
 

floquet/__init__.py

@@ -4,8 +4,8 @@
 from .floquet import (
     DisplacedState as DisplacedState,
     DisplacedStateFit as DisplacedStateFit,
-    DriveParameters as DriveParameters,
     FloquetAnalysis as FloquetAnalysis,
+    Model as Model,
 )
 from .options import Options as Options
 from .utils.file_io import (

floquet/amplitude_converters.py

@@ -16,16 +16,12 @@ class ChiacToAmp:
     """  # noqa E501
 
     def __init__(
-        self,
-        H0: qt.Qobj,
-        H1: qt.Qobj,
-        state_indices: list,
-        omega_d_linspace: np.ndarray,
+        self, H0: qt.Qobj, H1: qt.Qobj, state_indices: list, omega_d_values: np.ndarray
     ):
         self.H0 = H0
         self.H1 = H1
         self.state_indices = state_indices
-        self.omega_d_linspace = omega_d_linspace
+        self.omega_d_linspace = omega_d_values
 
     def amplitudes_for_omega_d(self, chi_ac_linspace: np.ndarray) -> np.ndarray:
         r"""Return drive amplitudes corresponding to $\chi_{\rm ac}$ values."""

floquet/displaced_state.py

@@ -7,7 +7,7 @@
 import qutip as qt
 import scipy as sp
 
-from .drive_parameters import DriveParameters
+from .model import Model
 from .options import Options
 from .utils.parallel import parallel_map
 
@@ -17,20 +17,17 @@ class DisplacedState:
 
     Parameters:
         hilbert_dim: Hilbert space dimension
-        drive_parameters: Drive parameters used
+        model: Model including the Hamiltonian, drive amplitudes, frequencies,
+            state indices
         state_indices: States of interest
         options: Options used
     """
 
     def __init__(
-        self,
-        hilbert_dim: int,
-        drive_parameters: DriveParameters,
-        state_indices: list,
-        options: Options,
+        self, hilbert_dim: int, model: Model, state_indices: list, options: Options
     ):
         self.hilbert_dim = hilbert_dim
-        self.drive_parameters = drive_parameters
+        self.model = model
         self.state_indices = state_indices
         self.options = options
         self.exponent_pair_idx_map = self._create_exponent_pair_idx_map()
@@ -73,19 +70,16 @@ def overlap_with_bare_states(
             def _compute_bare_state(
                 omega_d: float, _array_idx: int = array_idx, _state_idx: int = state_idx
             ) -> np.ndarray:
-                omega_d_idx = self.drive_parameters.omega_d_to_idx(omega_d)
+                omega_d_idx = self.model.omega_d_to_idx(omega_d)
                 return self.displaced_state(
                     omega_d,
-                    self.drive_parameters.drive_amplitudes[amp_idx_0, omega_d_idx],
+                    self.model.drive_amplitudes[amp_idx_0, omega_d_idx],
                     _state_idx,
                     coefficients=coefficients[_array_idx],
                 ).full()[:, 0]
 
             bare_states = np.array(
-                [
-                    _compute_bare_state(omega_d)
-                    for omega_d in self.drive_parameters.omega_d_values
-                ],
+                [_compute_bare_state(omega_d) for omega_d in self.model.omega_d_values],
                 dtype=complex,
             )
             # bare states may differ as a function of omega_d, hence the bare states
@@ -123,8 +117,8 @@ def _run_overlap_displaced(omega_d_amp: tuple[float, float]) -> np.ndarray:
             omega_d, amp = omega_d_amp
             for array_idx, state_idx in enumerate(self.state_indices):
                 floquet_mode_for_idx = floquet_modes[
-                    self.drive_parameters.omega_d_to_idx(omega_d),
-                    self.drive_parameters.amp_to_idx(amp, omega_d),
+                    self.model.omega_d_to_idx(omega_d),
+                    self.model.amp_to_idx(amp, omega_d),
                     array_idx,
                 ]
                 disp_state = self.displaced_state(
@@ -138,18 +132,16 @@ def _run_overlap_displaced(omega_d_amp: tuple[float, float]) -> np.ndarray:
                 )
             return overlap
 
-        omega_d_amp_params = self.drive_parameters.omega_d_amp_params(amp_idxs)
-        amp_range_vals = self.drive_parameters.drive_amplitudes[
-            amp_idxs[0] : amp_idxs[1]
-        ]
+        omega_d_amp_params = self.model.omega_d_amp_params(amp_idxs)
+        amp_range_vals = self.model.drive_amplitudes[amp_idxs[0] : amp_idxs[1]]
         result = list(
             parallel_map(
                 self.options.num_cpus, _run_overlap_displaced, omega_d_amp_params
             )
         )
         return np.array(result).reshape(
             (
-                len(self.drive_parameters.omega_d_values),
+                len(self.model.omega_d_values),
                 len(amp_range_vals),
                 len(self.state_indices),
             )
@@ -211,12 +203,8 @@ def _create_exponent_pair_idx_map(self) -> dict:
         but the fit is nominally set to order four. We additionally eliminate the
         constant term that should always be either zero or one.
         """
-        cutoff_omega_d = min(
-            len(self.drive_parameters.omega_d_values), self.options.fit_cutoff
-        )
-        cutoff_amp = min(
-            len(self.drive_parameters.drive_amplitudes), self.options.fit_cutoff
-        )
+        cutoff_omega_d = min(len(self.model.omega_d_values), self.options.fit_cutoff)
+        cutoff_amp = min(len(self.model.drive_amplitudes), self.options.fit_cutoff)
         idx_exp_map = [
             (idx_1, idx_2)
             for idx_1 in range(cutoff_omega_d)

floquet/floquet.py

@@ -6,7 +6,7 @@
 import qutip as qt
 
 from .displaced_state import DisplacedState, DisplacedStateFit
-from .drive_parameters import DriveParameters
+from .model import Model
 from .options import Options
 from .utils.file_io import Serializable
 from .utils.parallel import parallel_map
@@ -20,12 +20,8 @@ class FloquetAnalysis(Serializable):
     workflow, see the [transmon](../examples/transmon) tutorial.
 
     Arguments:
-        H0: Drift Hamiltonian, which must be diagonal and provided in units such that
-            H0 can be passed directly to qutip.
-        H1: Drive operator, which should be unitless (for instance the charge-number
-            operator n of the transmon). It will be multiplied by a drive amplitude
-            that we scan over from drive_parameters.drive_amplitudes.
-        drive_parameters: Class specifying the drive amplitudes and frequencies
+        model: Class specifying the model, including the Hamiltonian, drive amplitudes,
+            frequencies
         state_indices: State indices of interest. Defaults to [0, 1], indicating the two
             lowest-energy states.
         options: Options for the Floquet analysis.
@@ -34,37 +30,22 @@ class FloquetAnalysis(Serializable):
 
     def __init__(
         self,
-        H0: qt.Qobj | np.ndarray | list,
-        H1: qt.Qobj | np.ndarray | list,
-        drive_parameters: DriveParameters,
+        model: Model,
         state_indices: list | None = None,
         options: Options = Options(),  # noqa B008
         init_data_to_save: dict | None = None,
     ):
         if state_indices is None:
             state_indices = [0, 1]
-        if not isinstance(H0, qt.Qobj):
-            H0 = qt.Qobj(np.array(H0, dtype=complex))
-        if not isinstance(H1, qt.Qobj):
-            H1 = qt.Qobj(np.array(H1, dtype=complex))
-        self.H0 = H0
-        self.H1 = H1
-        self.drive_parameters = drive_parameters
+        self.model = model
         self.state_indices = state_indices
         self.options = options
         self.init_data_to_save = init_data_to_save
-        # Save in _init_attrs for later re-initialization. Everything added to self
-        # after this is a derived quantity
-        self.hilbert_dim = H0.shape[0]
+        self.hilbert_dim = model.H0.shape[0]
 
     def __str__(self) -> str:
         return "Running floquet simulation with parameters: \n" + super().__str__()
 
-    def hamiltonian(self, params: tuple[float, float]) -> list[qt.Qobj]:
-        """Return the Hamiltonian we actually simulate."""
-        omega_d, amp = params
-        return [self.H0, [amp * self.H1, lambda t, _: np.cos(omega_d * t)]]
-
     def run_one_floquet(
         self, omega_d_amp: tuple[float, float]
     ) -> tuple[np.ndarray, qt.Qobj]:
@@ -78,7 +59,7 @@ def run_one_floquet(
         omega_d, _ = omega_d_amp
         T = 2.0 * np.pi / omega_d
         f_modes_0, f_energies_0 = qt.floquet_modes(
-            self.hamiltonian(omega_d_amp),  # type: ignore
+            self.model.hamiltonian(omega_d_amp),  # type: ignore
             T,
             options=qt.Options(nsteps=self.options.nsteps),
         )
@@ -88,7 +69,7 @@ def run_one_floquet(
                 f_modes_0,
                 f_energies_0,
                 sampling_time,
-                self.hamiltonian(omega_d_amp),  # type: ignore
+                self.model.hamiltonian(omega_d_amp),  # type: ignore
                 T,
                 options=qt.Options(nsteps=self.options.nsteps),
             )
@@ -223,8 +204,8 @@ def run(self, filepath: str | None = None) -> dict:
 
         # initialize all arrays that will contain our data
         array_shape = (
-            len(self.drive_parameters.omega_d_values),
-            len(self.drive_parameters.drive_amplitudes),
+            len(self.model.omega_d_values),
+            len(self.model.drive_amplitudes),
             len(self.state_indices),
         )
         bare_state_overlaps = np.zeros(array_shape)
@@ -234,8 +215,8 @@ def run(self, filepath: str | None = None) -> dict:
         floquet_modes = np.zeros((*array_shape, self.hilbert_dim), dtype=complex)
         avg_excitation = np.zeros(
             (
-                len(self.drive_parameters.omega_d_values),
-                len(self.drive_parameters.drive_amplitudes),
+                len(self.model.omega_d_values),
+                len(self.model.drive_amplitudes),
                 self.hilbert_dim,
             )
         )
@@ -247,12 +228,11 @@ def run(self, filepath: str | None = None) -> dict:
         # coefficients, whereas for the Blais calculation, the bare modes are specified
         # as actual kets.
         prev_f_modes_arr = np.tile(
-            self.bare_state_array()[None, :, :],
-            (len(self.drive_parameters.omega_d_values), 1, 1),
+            self.bare_state_array()[None, :, :], (len(self.model.omega_d_values), 1, 1)
         )
         displaced_state = DisplacedStateFit(
             hilbert_dim=self.hilbert_dim,
-            drive_parameters=self.drive_parameters,
+            model=self.model,
             state_indices=self.state_indices,
             options=self.options,
         )
@@ -264,13 +244,13 @@ def run(self, filepath: str | None = None) -> dict:
         )
         num_fit_ranges = int(np.ceil(1 / self.options.fit_range_fraction))
         num_amp_pts_per_range = int(
-            np.floor(len(self.drive_parameters.drive_amplitudes) / num_fit_ranges)
+            np.floor(len(self.model.drive_amplitudes) / num_fit_ranges)
         )
         for amp_range_idx in range(num_fit_ranges):
             print(f"calculating for amp_range_idx={amp_range_idx}")
             # edge case if range doesn't fit in neatly
             if amp_range_idx == num_fit_ranges - 1:
-                amp_range_idx_final = len(self.drive_parameters.drive_amplitudes)
+                amp_range_idx_final = len(self.model.drive_amplitudes)
             else:
                 amp_range_idx_final = (amp_range_idx + 1) * num_amp_pts_per_range
             amp_idxs = [amp_range_idx * num_amp_pts_per_range, amp_range_idx_final]
@@ -305,7 +285,7 @@ def run(self, filepath: str | None = None) -> dict:
             ovlp_with_bare_states = displaced_state.overlap_with_bare_states(
                 amp_idxs[0], previous_coefficients, floquet_modes_for_range
             )
-            omega_d_amp_slice = list(self.drive_parameters.omega_d_amp_params(amp_idxs))
+            omega_d_amp_slice = list(self.model.omega_d_amp_params(amp_idxs))
             # Compute the fitted 'ideal' displaced state, excluding those
             # floquet modes experiencing resonances.
             new_coefficients = displaced_state.displaced_states_fit(
@@ -333,8 +313,8 @@ def run(self, filepath: str | None = None) -> dict:
         # (stored in intermediate_displaced_state_overlaps) to obtain the mask with
         # which we exclude some data from the fit (because we suspect they've hit
         # resonances).
-        amp_idxs = [0, len(self.drive_parameters.drive_amplitudes)]
-        omega_d_amp_slice = list(self.drive_parameters.omega_d_amp_params(amp_idxs))
+        amp_idxs = [0, len(self.model.drive_amplitudes)]
+        omega_d_amp_slice = list(self.model.omega_d_amp_params(amp_idxs))
         full_displaced_fit = displaced_state.displaced_states_fit(
             omega_d_amp_slice, intermediate_displaced_state_overlaps, floquet_modes
         )
@@ -371,14 +351,12 @@ def _floquet_main_for_amp_range(
         prev_f_modes_arr: np.ndarray,
     ) -> tuple:
         """Run the floquet simulation over a specific amplitude range."""
-        amp_range_vals = self.drive_parameters.drive_amplitudes[
-            amp_idxs[0] : amp_idxs[1]
-        ]
+        amp_range_vals = self.model.drive_amplitudes[amp_idxs[0] : amp_idxs[1]]
 
         def _run_floquet_and_calculate(
             omega_d: float,
         ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
-            omega_d_idx = self.drive_parameters.omega_d_to_idx(omega_d)
+            omega_d_idx = self.model.omega_d_to_idx(omega_d)
             amps_for_omega_d = amp_range_vals[:, omega_d_idx]
             avg_excitation_arr = np.zeros((len(amps_for_omega_d), self.hilbert_dim))
             quasienergies_arr = np.zeros_like(avg_excitation_arr)
@@ -414,7 +392,7 @@ def _run_floquet_and_calculate(
             parallel_map(
                 self.options.num_cpus,
                 _run_floquet_and_calculate,
-                self.drive_parameters.omega_d_values,
+                self.model.omega_d_values,
             )
         )
         (
@@ -425,32 +403,20 @@ def _run_floquet_and_calculate(
         ) = list(zip(*floquet_data, strict=True))
         floquet_mode_array = np.array(all_modes_quasies_ovlps, dtype=complex).reshape(
             (
-                len(self.drive_parameters.omega_d_values),
+                len(self.model.omega_d_values),
                 len(amp_range_vals),
                 len(self.state_indices),
                 1 + self.hilbert_dim,
             )
         )
         f_modes_last_amp = np.array(f_modes_last_amp, dtype=complex).reshape(
-            (
-                len(self.drive_parameters.omega_d_values),
-                self.hilbert_dim,
-                self.hilbert_dim,
-            )
+            (len(self.model.omega_d_values), self.hilbert_dim, self.hilbert_dim)
         )
         all_avg_excitation = np.array(all_avg_excitation).reshape(
-            (
-                len(self.drive_parameters.omega_d_values),
-                len(amp_range_vals),
-                self.hilbert_dim,
-            )
+            (len(self.model.omega_d_values), len(amp_range_vals), self.hilbert_dim)
         )
         all_quasienergies = np.array(all_quasienergies).reshape(
-            (
-                len(self.drive_parameters.omega_d_values),
-                len(amp_range_vals),
-                self.hilbert_dim,
-            )
+            (len(self.model.omega_d_values), len(amp_range_vals), self.hilbert_dim)
         )
         bare_state_overlaps = np.abs(floquet_mode_array[..., 0])
         floquet_modes = floquet_mode_array[..., 1:]