from c3.c3objs import C3obj, hjson_encode
from c3.c3objs import Quantity as Qty
from c3.libraries.envelopes import envelopes
from c3.utils.tf_utils import tf_complexify
import tensorflow as tf
import numpy as np
import types
import hjson
from typing import Callable, Union, Dict
components = dict()
[docs]def comp_reg_deco(func):
"""
Decorator for making registry of functions
"""
components[str(func.__name__)] = func
return func
[docs]@comp_reg_deco
class Envelope(C3obj):
"""
Represents the envelopes shaping a pulse.
Parameters
----------
shape: Callable
function evaluating the shape in time
"""
def __init__(
self,
name: str,
desc: str = " ",
comment: str = " ",
params: Dict[str, Qty] = {},
shape: Union[Callable, str] = None,
use_t_before=False,
):
if isinstance(shape, str):
self.shape = envelopes[shape]
else:
self.shape = shape
default_params = {
"amp": Qty(value=1.0, min_val=-1.0, max_val=+1.5, unit="V"),
"delta": Qty(value=0.0, min_val=-5.0, max_val=+5.0, unit="V"),
"freq_offset": Qty(value=0.0, min_val=-1.0, max_val=+1.0, unit="Hz 2pi"),
"xy_angle": Qty(value=0.0, min_val=-1.0, max_val=+1.0, unit="rad"),
"sigma": Qty(value=5e-9, min_val=-2.0, max_val=+2.0, unit="s"),
"t_final": Qty(value=1.0, min_val=-1.0, max_val=+1.0, unit="s"),
}
default_params.update(params)
self.set_use_t_before(use_t_before)
super().__init__(
name=name,
desc=desc,
comment=comment,
params=default_params,
)
[docs] def write_config(self, filepath: str) -> None:
"""
Write dictionary to a HJSON file.
"""
with open(filepath, "w") as cfg_file:
hjson.dump(self.asdict(), cfg_file, default=hjson_encode)
[docs] def asdict(self) -> dict:
params = {}
for key, item in self.params.items():
params[key] = item.asdict()
return {
"name": self.name,
"c3type": self.__class__.__name__,
"shape": self.shape.__name__,
"params": params,
}
[docs] def __str__(self) -> str:
return hjson.dumps(self.asdict(), default=hjson_encode)
[docs] def __repr__(self) -> str:
repr_str = self.name + ":: "
for key, item in self.params.items():
repr_str += str(key) + " : " + str(item) + ", "
repr_str += "shape: " + self.shape.__name__ + ", "
return repr_str
[docs] def set_use_t_before(self, use_t_before):
if use_t_before:
self.get_shape_values = self._get_shape_values_before
else:
self.get_shape_values = self._get_shape_values_just
[docs] def compute_mask(self, ts, t_end) -> tf.Tensor:
"""Compute a mask to cut out a signal after t_final.
Parameters
----------
ts : tf.Tensor
Vector of time steps.
Returns
-------
tf.Tensor
[description]
"""
t_final = tf.minimum(self.params["t_final"].get_value(), t_end)
dt = ts[1] - ts[0]
return tf_complexify(
tf.sigmoid((ts / dt + 0.001) * 1e6)
* tf.sigmoid((0.999 * t_final - ts) / dt * 1e6)
)
def _get_shape_values_before(self, ts, t_final=1):
"""Return the value of the shape function at the specified times. With the offset, we make sure the
signal starts with amplitude zero by subtracting the shape value at time -dt.
Parameters
----------
ts : tf.Tensor
Vector of time samples.
"""
t_before = 2 * ts[0] - ts[1] # t[0] - (t[1] - t[0])
offset = self.shape(t_before, self.params)
mask = self.compute_mask(ts, t_final)
return mask * (self.shape(ts, self.params) - offset)
def _get_shape_values_just(self, ts, t_final=1):
"""Return the value of the shape function at the specified times.
Parameters
----------
ts : tf.Tensor
Vector of time samples.
"""
mask = self.compute_mask(ts, t_final)
return mask * self.shape(ts, self.params)
[docs]@comp_reg_deco
class EnvelopeDrag(Envelope):
def __init__(
self,
name: str,
desc: str = " ",
comment: str = " ",
params: dict = {},
shape: types.FunctionType = None,
use_t_before=False,
):
super().__init__(
name=name,
desc=desc,
comment=comment,
params=params,
shape=shape,
use_t_before=use_t_before,
)
self.set_use_t_before(use_t_before)
[docs] def set_use_t_before(self, use_t_before):
if use_t_before:
self.base_env = super()._get_shape_values_before
else:
self.base_env = super()._get_shape_values_just
[docs] def get_shape_values(self, ts, t_final=1):
dt = ts[1] - ts[0]
with tf.GradientTape() as t:
t.watch(ts)
env = tf.math.real(self.base_env(ts, t_final))
denv = (
t.gradient(env, ts, unconnected_gradients=tf.UnconnectedGradients.ZERO) * dt
) # Derivative W.R.T. to bins
delta = self.params["delta"].get_value()
return tf.complex(env, -denv * delta)
[docs]@comp_reg_deco
class EnvelopeNetZero(Envelope):
"""
Represents the envelopes shaping a pulse.
Parameters
----------
shape: function
function evaluating the shape in time
params: dict
Parameters of the envelope
Note: t_final
"""
def __init__(
self,
name: str,
desc: str = " ",
comment: str = " ",
params: dict = {},
shape: types.FunctionType = None,
use_t_before=False,
):
super().__init__(
name=name,
desc=desc,
comment=comment,
params=params,
shape=shape,
use_t_before=use_t_before,
)
self.set_use_t_before(use_t_before)
[docs] def set_use_t_before(self, use_t_before):
if use_t_before:
self.base_env = super()._get_shape_values_before
else:
self.base_env = super()._get_shape_values_just
[docs] def get_shape_values(self, ts):
"""Return the value of the shape function at the specified times.
Parameters
----------
ts : tf.Tensor
Vector of time samples.
t_before : tf.float64
Offset the beginning of the shape by this time.
"""
N_red = len(ts) // 2
ts_red = tf.split(ts, [N_red, len(ts) - N_red], 0)[0]
shape_values = self.base_env(ts=ts_red)
netzero_shape_values = tf.concat(
[shape_values, -shape_values, [0] * (len(ts) % 2)], axis=0
)
return netzero_shape_values
[docs]@comp_reg_deco
class Carrier(C3obj):
"""Represents the carrier of a pulse."""
def __init__(
self,
name: str,
desc: str = " ",
comment: str = " ",
params: dict = {},
):
params_default = {
"freq": Qty(value=0.0, min_val=-1.0, max_val=+1.0, unit="V"),
"framechange": Qty(value=0.0, min_val=-np.pi, max_val=np.pi, unit="rad"),
}
params_default.update(params)
super().__init__(
name=name,
desc=desc,
comment=comment,
params=params_default,
)
[docs] def write_config(self, filepath: str) -> None:
"""
Write dictionary to a HJSON file.
"""
with open(filepath, "w") as cfg_file:
hjson.dump(self.asdict(), cfg_file, default=hjson_encode)
[docs] def __repr__(self) -> str:
repr_str = self.name + ":: "
for key, item in self.params.items():
repr_str += str(key) + " : " + str(item) + ", "
return repr_str