[Feature, Testing] Allow nested calls to ConcretizedCallable, more tests

pyproject.toml

@@ -117,7 +117,7 @@ serve = "mkdocs serve --dev-addr localhost:8000"
 
 [tool.ruff]
 lint.select = ["E", "F", "I", "Q"]
-lint.extend-ignore = ["F841"]
+lint.extend-ignore = ["F841", "E731"]
 line-length = 100
 
 [tool.ruff.lint.isort]

pyqtorch/__init__.py

@@ -55,7 +55,16 @@
     Scale,
     Sequence,
 )
-from .embed import ConcretizedCallable, Embedding
+from .embed import (
+    ConcretizedCallable,
+    Embedding,
+    cos,
+    log,
+    sin,
+    sqrt,
+    tan,
+    tanh,
+)
 from .hamiltonians import HamiltonianEvolution, Observable
 from .noise import (
     AmplitudeDamping,

pyqtorch/composite/compose.py

@@ -10,7 +10,7 @@
 from torch.nn import Module, ModuleList, ParameterDict
 
 from pyqtorch.apply import apply_operator
-from pyqtorch.embed import Embedding
+from pyqtorch.embed import ConcretizedCallable, Embedding
 from pyqtorch.matrices import add_batch_dim
 from pyqtorch.primitives import CNOT, RX, RY, Parametric, Primitive
 from pyqtorch.utils import (
@@ -35,7 +35,9 @@ class Scale(Sequence):
     """
 
     def __init__(
-        self, operations: Union[Primitive, Sequence, Add], param_name: str | Tensor
+        self,
+        operations: Union[Primitive, Sequence, Add],
+        param_name: str | Tensor,
     ):
         """
         Initializes a Scale object.
@@ -69,12 +71,14 @@ def forward(
         if embedding is not None:
             values = embedding(values)
 
-        scale = (
-            values[self.param_name]
-            if isinstance(self.param_name, str)
-            else self.param_name
-        )
-        return scale * self.operations[0].forward(state, values, embedding)
+        if isinstance(self.param_name, str):
+            scale = values[self.param_name]
+        elif isinstance(self.param_name, Tensor):
+            scale = self.param_name
+        elif isinstance(self.param_name, ConcretizedCallable):
+            scale = self.param_name(values)
+
+        return scale * self.operations[0].forward(state, values)
 
     def tensor(
         self,
@@ -97,12 +101,14 @@ def tensor(
         if embedding is not None:
             values = embedding(values)
 
-        scale = (
-            values[self.param_name]
-            if isinstance(self.param_name, str)
-            else self.param_name
-        )
-        return scale * self.operations[0].tensor(values, embedding, full_support)
+        if isinstance(self.param_name, str):
+            scale = values[self.param_name]
+        elif isinstance(self.param_name, (Tensor, int, float)):
+            scale = self.param_name
+        elif isinstance(self.param_name, ConcretizedCallable):
+            scale = self.param_name(values)
+
+        return scale * self.operations[0].tensor(values, full_support=full_support)
 
     def flatten(self) -> list[Scale]:
         """This method should only be called in the AdjointExpectation,

pyqtorch/embed.py

@@ -5,6 +5,7 @@
 from typing import Any, Tuple
 
 from numpy.typing import ArrayLike, DTypeLike
+from torch import Tensor
 
 logger = getLogger(__name__)
 
@@ -21,7 +22,7 @@
     "sub": ("jax.numpy", "subtract"),
     "div": ("jax.numpy", "divide"),
 }
-DEFAULT_TORCH_MAPPING: dict = dict()
+DEFAULT_TORCH_MAPPING = {"hs": ("pyqtorch.utils", "heaviside")}
 DEFAULT_NUMPY_MAPPING = {
     "mul": ("numpy", "multiply"),
     "sub": ("numpy", "subtract"),
@@ -74,8 +75,8 @@ class ConcretizedCallable:
 
     def __init__(
         self,
-        call_name: str,
-        abstract_args: list[str | float | int],
+        call_name: str = "",
+        abstract_args: list[str | float | int | complex | ConcretizedCallable] = ["x"],
         instruction_mapping: dict[str, Tuple[str, str]] = dict(),
         engine_name: str = "torch",
         device: str = "cpu",
@@ -113,17 +114,94 @@ def __init__(
     def evaluate(self, inputs: dict[str, ArrayLike] = dict()) -> ArrayLike:
         arraylike_args = []
         for symbol_or_numeric in self.abstract_args:
-            if isinstance(symbol_or_numeric, (float, int)):
+            if isinstance(symbol_or_numeric, ConcretizedCallable):
+                arraylike_args.append(symbol_or_numeric(inputs))
+            if isinstance(symbol_or_numeric, (float, int, Tensor)):
                 arraylike_args.append(
                     self.arraylike_fn(symbol_or_numeric, device=self.device)
                 )
             elif isinstance(symbol_or_numeric, str):
                 arraylike_args.append(inputs[symbol_or_numeric])
         return self.engine_call(*arraylike_args)  # type: ignore[misc]
 
+    @classmethod
+    def _get_independent_args(cls, cc: ConcretizedCallable) -> set:
+        out: set = set()
+        if len(cc.abstract_args) == 1 and isinstance(cc.abstract_args[0], str):
+            return set([cc.abstract_args[0]])
+        else:
+            for arg in cc.abstract_args:
+                if isinstance(arg, ConcretizedCallable):
+                    res = cls._get_independent_args(arg)
+                    out = out.union(res)
+                else:
+                    if isinstance(arg, str):
+                        out.add(arg)
+        return out
+
+    @property
+    def independent_args(self) -> list:
+        return list(self._get_independent_args(self))
+
     def __call__(self, inputs: dict[str, ArrayLike] = dict()) -> ArrayLike:
         return self.evaluate(inputs)
 
+    def __mul__(
+        self, other: str | int | float | complex | ConcretizedCallable
+    ) -> ConcretizedCallable:
+        return ConcretizedCallable("mul", [self, other])
+
+    def __rmul__(
+        self, other: str | int | float | ConcretizedCallable
+    ) -> ConcretizedCallable:
+        return ConcretizedCallable("mul", [other, self])
+
+    def __add__(
+        self, other: str | int | float | ConcretizedCallable
+    ) -> ConcretizedCallable:
+        return ConcretizedCallable("add", [self, other])
+
+    def __radd__(
+        self, other: str | int | float | ConcretizedCallable
+    ) -> ConcretizedCallable:
+        return ConcretizedCallable("add", [other, self])
+
+    def __sub__(
+        self, other: str | int | float | ConcretizedCallable
+    ) -> ConcretizedCallable:
+        return ConcretizedCallable("sub", [self, other])
+
+    def __rsub__(
+        self, other: str | int | float | ConcretizedCallable
+    ) -> ConcretizedCallable:
+        return ConcretizedCallable("sub", [other, self])
+
+    def __pow__(
+        self, other: str | int | float | ConcretizedCallable
+    ) -> ConcretizedCallable:
+        return ConcretizedCallable("pow", [self, other])
+
+    def __rpow__(
+        self, other: str | int | float | ConcretizedCallable
+    ) -> ConcretizedCallable:
+        return ConcretizedCallable("pow", [other, self])
+
+    def __truediv__(
+        self, other: str | int | float | ConcretizedCallable
+    ) -> ConcretizedCallable:
+        return ConcretizedCallable("div", [self, other])
+
+    def __rtruediv__(
+        self, other: str | int | float | ConcretizedCallable
+    ) -> ConcretizedCallable:
+        return ConcretizedCallable("div", [other, self])
+
+    def __repr__(self) -> str:
+        return f"{self.call_name}({self.abstract_args})"
+
+    def __neg__(self) -> ConcretizedCallable:
+        return -1 * self
+
     @property
     def device(self) -> str:
         return self._device
@@ -150,6 +228,30 @@ def init_param(
         return engine.random.uniform(0, 1)
 
 
+def sin(x: str | ConcretizedCallable):
+    return ConcretizedCallable("sin", [x])
+
+
+def cos(x: str | ConcretizedCallable):
+    return ConcretizedCallable("cos", [x])
+
+
+def log(x: str | ConcretizedCallable):
+    return ConcretizedCallable("log", [x])
+
+
+def tan(x: str | ConcretizedCallable):
+    return ConcretizedCallable("tan", [x])
+
+
+def tanh(x: str | ConcretizedCallable):
+    return ConcretizedCallable("tanh", [x])
+
+
+def sqrt(x: str | ConcretizedCallable):
+    return ConcretizedCallable("sqrt", [x])
+
+
 class Embedding:
     """A class relating variational and feature parameters used in ConcretizedCallable instances to
     parameter names used in gates.

pyqtorch/hamiltonians/evolution.py

@@ -12,7 +12,8 @@
 
 from pyqtorch.apply import apply_operator
 from pyqtorch.circuit import Sequence
-from pyqtorch.embed import Embedding
+from pyqtorch.composite import Scale
+from pyqtorch.embed import ConcretizedCallable, Embedding
 from pyqtorch.primitives import Primitive
 from pyqtorch.quantum_operation import QuantumOperation
 from pyqtorch.time_dependent.sesolve import sesolve
@@ -142,9 +143,10 @@ class HamiltonianEvolution(Sequence):
     def __init__(
         self,
         generator: TGenerator,
-        time: Tensor | str,
+        time: Tensor | str | ConcretizedCallable,
         qubit_support: Tuple[int, ...] | None = None,
         cache_length: int = 1,
+        duration: float | Tensor = 1.0,
         steps: int = 100,
         solver=SolverType.DP5_SE,
     ):
@@ -160,6 +162,19 @@ def __init__(
 
         self.solver_type = solver
         self.steps = steps
+        self.duration = duration
+        self.is_time_dependent = None
+
+        if (
+            isinstance(time, str)
+            or isinstance(time, Tensor)
+            or isinstance(time, ConcretizedCallable)
+        ):
+            self.time = time
+        else:
+            raise ValueError("time should be passed as str or tensor.")
+
+        self.has_time_param = self._has_time_param(generator)
 
         if isinstance(generator, Tensor):
             if qubit_support is None:
@@ -185,6 +200,7 @@ def __init__(
                     "Taking support from generator and ignoring qubit_support input."
                 )
             qubit_support = generator.qubit_support
+
             if is_parametric(generator):
                 generator = [generator]
                 self.generator_type = GeneratorType.PARAMETRIC_OPERATION
@@ -205,11 +221,6 @@ def __init__(
         super().__init__(generator)
         self._qubit_support = qubit_support  # type: ignore
 
-        if isinstance(time, str) or isinstance(time, Tensor):
-            self.time = time
-        else:
-            raise ValueError("time should be passed as str or tensor.")
-
         logger.debug("Hamiltonian Evolution initialized")
         if logger.isEnabledFor(logging.DEBUG):
             # When Debugging let's add logging and NVTX markers
@@ -246,6 +257,24 @@ def flatten(self) -> ModuleList:
     def param_name(self) -> Tensor | str:
         return self.time
 
+    def _has_time_param(self, generator: TGenerator) -> bool:
+        from pyqtorch.primitives import Parametric
+
+        res = False
+        if isinstance(self.time, Tensor):
+            return res
+        else:
+            if isinstance(generator, (Sequence, QuantumOperation)):
+                for m in generator.modules():
+                    if isinstance(m, (Scale, Parametric)):
+                        if self.time in getattr(m.param_name, "independent_args", []):
+                            # param_name is a ConcretizedCallable object
+                            res = True
+                        elif m.param_name == self.time:
+                            # param_name is a string
+                            res = True
+        return res
+
     def _symbolic_generator(
         self,
         values: dict,
@@ -341,20 +370,27 @@ def _forward_time(
     ) -> State:
         n_qubits = len(state.shape) - 1
         batch_size = state.shape[-1]
-        t_grid = torch.linspace(0, float(self.time), self.steps)
+        t_grid = torch.linspace(0, float(self.duration), self.steps)
 
-        values.update({embedding.tparam_name: torch.tensor(0.0)})  # type: ignore [dict-item]
-        embedded_params = embedding(values)
+        if embedding is not None:
+            values.update({embedding.tparam_name: torch.tensor(0.0)})  # type: ignore [dict-item]
+            embedded_params = embedding(values)
+        else:
+            embedded_params = values
 
         def Ht(t: torch.Tensor) -> torch.Tensor:
             """Accepts a value 't' for time and returns
             a (2**n_qubits, 2**n_qubits) Hamiltonian evaluated at time 't'.
             """
-            # We use the origial embedded params and return a new dict
+            # We use the original embedded params and return a new dict
             # where we reembedded all parameters depending on time with value 't'
-            reembedded_time_values = embedding.reembed_tparam(
-                embedded_params, torch.as_tensor(t)
-            )
+            if embedding is not None:
+                reembedded_time_values = embedding.reembed_tparam(
+                    embedded_params, torch.as_tensor(t)
+                )
+            else:
+                values[self.time] = torch.as_tensor(t)
+                reembedded_time_values = values
             return (
                 self.generator[0].tensor(reembedded_time_values, embedding).squeeze(2)
             )
@@ -388,8 +424,10 @@ def forward(
         Returns:
             The transformed state.
         """
-        if embedding is not None and getattr(embedding, "tparam_name", None):
-            return self._forward_time(state, values, embedding)
+        if self.has_time_param or (
+            embedding is not None and getattr(embedding, "tparam_name", None)
+        ):
+            return self._forward_time(state, values, embedding)  # type: ignore [arg-type]
 
         else:
             return self._forward(state, values, embedding)
@@ -420,10 +458,14 @@ def tensor(
             evolved_op = self._cache_hamiltonian_evo[values_cache_key]
         else:
             hamiltonian: torch.Tensor = self.create_hamiltonian(values, embedding)  # type: ignore [call-arg]
-            time_evolution: torch.Tensor = (
-                values[self.time] if isinstance(self.time, str) else self.time
-            )  # If `self.time` is a string / hence, a Parameter,
-            # we expect the user to pass it in the `values` dict
+
+            if isinstance(self.time, str):
+                time_evolution = values[self.time]
+            elif isinstance(self.time, ConcretizedCallable):
+                time_evolution = self.time(values)
+            else:
+                time_evolution = self.time
+
             evolved_op = evolve(hamiltonian, time_evolution)
             nb_cached = len(self._cache_hamiltonian_evo)