Ansatz¶

`Ansatz`	A base class for ansatz.
`ArrayAnsatz`	The ansatz of the Fock-Bargmann representation.
`PolyExpAnsatz`	This class represents the ansatz of a polynomial exponential function.

Inheritance Diagram¶

Ansatz¶

class mrmustard.physics.ansatz.Ansatz[source]¶

A base class for ansatz.

abstract property batch_dims: int¶: The number of batch dimensions of the ansatz.

abstract property batch_shape: tuple[int, ...]¶: The batch shape of the ansatz.

abstract property batch_size: int¶: The batch size of the ansatz.

abstract property conj: Ansatz¶: The conjugate of the ansatz.

abstract property core_dims: int¶: The number of core dimensions of the ansatz.

abstract property data: tuple | Tensor¶: The data of the ansatz. For now, it’s the triple for PolyExpAnsatz and the array for ArrayAnsatz.

abstract property num_vars: int¶: The number of variables of this ansatz.

abstract property scalar: Scalar¶: The scalar part of the ansatz. For now it’s c for PolyExpAnsatz and the array for ArrayAnsatz.

abstract classmethod from_dict(data)[source]¶

Deserialize an Ansatz.

Parameters:: data (dict[str, ArrayLike]) – The data to deserialize.
Returns:: An Ansatz.
Return type:: Ansatz

abstract concat(other, axis=0)[source]¶

Concatenates this ansatz with another along a specified batch axis.

All batch axes except the concatenation axis must agree in size. Core dimensions and other properties must match according to the specific ansatz type.

Parameters:

other (Ansatz) – Another ansatz of the same type.
axis (int) – The batch axis along which to concatenate (default 0).

Returns:

A new ansatz with concatenated batch dimensions.

Raises:

ValueError – If the ansatze have incompatible shapes or properties.

Return type:

Ansatz

abstract contract(other, idxs)[source]¶

Contract two ansatz along the specified core indices, broadcasting batch dimensions.

Parameters:

other (Ansatz) – Another ansatz.
idxs (tuple[Sequence[int], Sequence[int]]) – Tuple (idx_self, idx_other) of sequences of core-axis indices (0-based, relative to the core variables of each operand) to contract. The two sequences must have the same length. Negative indices are supported and are interpreted relative to the number of core variables.

Returns:

The resulting contracted ansatz.

Return type:

Ansatz

abstract reorder(order)[source]¶

Reorders the ansatz indices.

Parameters:: order (tuple[int, ...] | list[int]) – The desired order of the ansatz indices.
Returns:: A new Ansatz with reordered indices.
Return type:: Ansatz

abstract reorder_batch(order)[source]¶

Reorders the batch dimensions of the ansatz. The length of order must equal the number of batch dimensions. This method returns a new ansatz object.

Parameters:: order (Sequence[int]) – The desired order of the batch dimensions.
Returns:: A new Ansatz with reordered batch dimensions.
Return type:: Ansatz

abstract to_dict()[source]¶

Serialize an Ansatz.

Returns:: A dictionary containing the serialized data.
Return type:: dict[str, ArrayLike]

abstract trace(idx_z, idx_zconj)[source]¶

Implements the partial trace over the given index pairs.

Parameters:

idx_z (tuple[int, ...]) – The first part of the pairs of indices to trace over.
idx_zconj (tuple[int, ...]) – The second part.

Returns:

The traced-over ansatz.

Return type:

Ansatz

ArrayAnsatz¶

class mrmustard.physics.ansatz.ArrayAnsatz(array, batch_dims=0)[source]¶

Bases: Ansatz

The ansatz of the Fock-Bargmann representation.

Represents the ansatz as a multidimensional array.

Parameters:

array (Tensor) – A (potentially) batched array.
batch_dims (int) – The number of batch dimensions.

>>> import numpy as np
>>> from mrmustard.physics.ansatz import ArrayAnsatz
>>> array = np.random.random((2, 4, 5))
>>> ansatz = ArrayAnsatz(array)

property array: Tensor¶: The array of this ansatz.

property batch_dims: int¶: The number of batch dimensions of the ansatz.

property batch_shape: tuple[int, ...]¶: The batch shape of the ansatz.

property batch_size¶: The batch size of the ansatz.

property conj: ArrayAnsatz¶: The conjugate of the ansatz.

property core_dims: int¶: The number of core dimensions of this ansatz.

property core_shape: tuple[int, ...]¶: The core dimensions of this ansatz.

property data: Tensor¶: The data of the ansatz. For now, it’s the triple for PolyExpAnsatz and the array for ArrayAnsatz.

property num_vars: int¶: The number of variables of this ansatz.

property scalar: Scalar¶: The scalar part of the ansatz. I.e. the vacuum component of the Fock array, whatever it may be.

classmethod from_dict(data)[source]¶

Deserialize an Ansatz.

Parameters:: data (dict[str, ArrayLike]) – The data to deserialize.
Returns:: An Ansatz.
Return type:: ArrayAnsatz

concat(other, axis=0)[source]¶

Concatenates two ArrayAnsatz objects along the specified batch axis.

All batch axes except the concatenation axis must agree in size. Core shapes must also match.

Parameters:

other (Ansatz) – The other ArrayAnsatz to concatenate with.
axis (int) – The batch axis along which to concatenate (default 0).

Returns:

A new ArrayAnsatz with the concatenated batch dimensions.

Raises:

TypeError – If other is not an ArrayAnsatz.
ValueError – If batch dimensions don’t match (except at axis) or if core shapes don’t match.

Return type:

ArrayAnsatz

>>> from mrmustard.physics.ansatz import ArrayAnsatz
>>> import numpy as np
>>> array1 = np.random.random((2, 4, 5))
>>> ansatz1 = ArrayAnsatz(array1, batch_dims=1)
>>> array2 = np.random.random((3, 4, 5))
>>> ansatz2 = ArrayAnsatz(array2, batch_dims=1)
>>> concatenated = ansatz1.concat(ansatz2, axis=0)
>>> assert concatenated.batch_shape == (5,)
>>> assert concatenated.core_shape == (4, 5)

contract(other, idxs)[source]¶

Contract along specified core axes, broadcasting batch dimensions.

Parameters:

other (Ansatz) – The other ArrayAnsatz to contract with.
idxs (tuple[Sequence[int], Sequence[int]]) – Tuple (idx_self, idx_other) of core-axis indices to contract (0-based relative to core dims). Negative indices allowed.

Returns:

The contracted ArrayAnsatz, with kept core axes ordered as [self non-contracted] + [other non-contracted]. Batch dims are broadcast and preserved.

Raises:

TypeError – If other is not an ArrayAnsatz.
ValueError – If index sequences have incorrect length, duplicates, or out-of-range indices.

Return type:

ArrayAnsatz

>>> from mrmustard.physics.ansatz import ArrayAnsatz
>>> from mrmustard import math
>>> array1 = math.arange(20).reshape((1, 4, 5))
>>> array2 = math.arange(72).reshape((3, 4, 6))
>>> ansatz1 = ArrayAnsatz(array1, batch_dims=1)
>>> ansatz2 = ArrayAnsatz(array2, batch_dims=1)
>>> # broadcast 1 and 3 while contracting 4 and 4, leaving 3 (batch), 5 and 6:
>>> contracted = ansatz1.contract(ansatz2, ([0], [0]))
>>> assert contracted.array.shape == (3, 5, 6)

reduce(shape)[source]¶

Returns a new ArrayAnsatz with a sliced core shape.

Parameters:: shape (Sequence[int]) – The shape of the array of the returned ArrayAnsatz.
Return type:: ArrayAnsatz

>>> from mrmustard import math
>>> from mrmustard.physics.ansatz import ArrayAnsatz
>>> array1 = math.arange(27).reshape((3, 3, 3))
>>> fock1 = ArrayAnsatz(array1)
>>> fock2 = fock1.reduce((3, 3, 3))
>>> assert fock1 == fock2
>>> fock3 = fock1.reduce((2, 2, 2))
>>> array3 = math.astensor([[[0, 1], [3, 4]], [[9, 10], [12, 13]]])
>>> assert fock3 == ArrayAnsatz(array3)
>>> fock4 = fock1.reduce((1, 3, 1))
>>> array4 = math.astensor([[[0], [3], [6]]])
>>> assert fock4 == ArrayAnsatz(array4)

reorder(order)[source]¶

Reorders the ansatz indices.

Parameters:: order (tuple[int, ...] | list[int]) – The desired order of the ansatz indices.
Returns:: A new Ansatz with reordered indices.
Return type:: ArrayAnsatz

reorder_batch(order)[source]¶

Reorders the batch dimensions of the ansatz. The length of order must equal the number of batch dimensions. This method returns a new ansatz object.

Parameters:: order (Sequence[int]) – The desired order of the batch dimensions.
Returns:: A new Ansatz with reordered batch dimensions.

to_dict()[source]¶

Serialize an Ansatz.

Returns:: A dictionary containing the serialized data.
Return type:: dict[str, ArrayLike]

trace(idx_z, idx_zconj)[source]¶

Implements the partial trace over the given index pairs.

Parameters:

idx_z (tuple[int, ...]) – The first part of the pairs of indices to trace over.
idx_zconj (tuple[int, ...]) – The second part.

Returns:

The traced-over ansatz.

Return type:

ArrayAnsatz

PolyExpAnsatz¶

class mrmustard.physics.ansatz.PolyExpAnsatz(A, b, c, name='', lin_sup=False)[source]¶

Bases: Ansatz

This class represents the ansatz of a polynomial exponential function.

Namely,

\(F^{(i)}(z) = \sum_k c^{(i)}_{k} \partial_y^k \textrm{exp}(\frac{1}{2}(z,y)^T A^{(i)} (z,y) + (z,y)^T b^{(i)})|_{y=0}\)

with k and i multi-indices. The i multi-index is a batch index of shape L that can be used for linear superposition or batching purposes. Each of the c^{(i)}_k tensors are contracted with the array of derivatives \(\partial_y^k\) to form polynomials of derivatives. The tensors \(c^{(i)}_{k}\) contain the coefficients of the polynomial of derivatives and have shape (*L, *derived), where *derived is the shape of the derived variables, which implies len(c.shape[1:]) = m. The matrices \(A^{(i)}\) and vectors \(b^{(i)}\) are the parameters of the exponential terms in the ansatz, with \(z\in\mathbb{C}^{n}\) and \(y\in\mathbb{C}^{m}\). A and b have shape (*L, n+m, n+m) and (*L, n+m), respectively.

Parameters:

A (Matrix) – A batch of quadratic coefficient \(A^{(i)}\).
b (Vector) – A batch of linear coefficients \(b^{(i)}\).
c (Tensor) – A batch of arrays \(c^{(i)}\).
name (str) – The name of the ansatz.
lin_sup (bool) – Whether to include linear superposition axes in the batch dimensions.

>>> from mrmustard.physics.ansatz import PolyExpAnsatz
>>> import numpy as np
>>> A = np.random.random((3,3))  # no batch
>>> b = np.random.random((3,))
>>> c = np.random.random()
>>> F = PolyExpAnsatz(A, b, c)
>>> assert F(1.0, 2.0, 3.0).shape == ()
>>> A = np.random.random((10,3,3))  # batch of 10
>>> b = np.random.random((10,3))
>>> c = np.random.random((10,))
>>> F = PolyExpAnsatz(A, b, c)
>>> assert F(1.0, 2.0, 3.0).shape == (10,)
>>> A = np.random.random((10,3,3))  # batch of 10
>>> b = np.random.random((10,3))
>>> c = np.random.random((10,7))
>>> F = PolyExpAnsatz(A, b, c)
>>> assert F(1.0, 2.0).shape == (10,)  # two CV variables, one derived
>>> A = np.random.random((10,3,3))  # batch of 10
>>> b = np.random.random((10,3))
>>> c = np.random.random((10,7,5))
>>> F = PolyExpAnsatz(A, b, c)
>>> assert F(1.0).shape == (10,)  # one CV variable, two derived
>>> assert F([1.0, 2.0, 3.0]).shape == (3,10)  # batch of 3 inputs

property A: ComplexMatrix¶: The batch of quadratic coefficient \(A^{(i)}\).

property b: ComplexVector¶: The batch of linear coefficients \(b^{(i)}\).

property batch_dims: int¶: The number of batch dimensions of the ansatz.

property batch_shape: tuple[int, ...]¶: The batch shape of the ansatz.

property batch_size: int¶: The batch size of the ansatz.

property c: ComplexTensor¶: The batch of polynomial coefficients \(c^{(i)}_{k}\).

property conj: PolyExpAnsatz¶: The conjugate of the ansatz.

property core_dims: int¶: The number of core variables of the ansatz. Equivalent to self.num_CV_vars.

property data: tuple[mrmustard.utils.typing.ComplexMatrix, mrmustard.utils.typing.ComplexVector, mrmustard.utils.typing.ComplexTensor]¶: Returns the triple, which is necessary to reinstantiate the ansatz.

property num_CV_vars: int¶: The number of continuous variables that remain after the polynomial of derivatives is applied. This is the number of continuous variables of the Ansatz function itself.

property num_derived_vars: int¶: The number of derived variables that are derived by the polynomial of derivatives.

property num_vars¶: The number of variables of this ansatz.

property PS: PolyExpAnsatz¶: The ansatz defined using real (i.e., phase-space) variables.

property scalar: Scalar¶: The scalar part of the ansatz, i.e. F(0).

property shape_derived_vars: tuple[int, ...]¶: The shape of the coefficients of the polynomial of derivatives.

property triple: tuple[mrmustard.utils.typing.ComplexMatrix, mrmustard.utils.typing.ComplexVector, mrmustard.utils.typing.ComplexTensor]¶: Returns the triple of parameters of the exponential part of the ansatz.

classmethod from_dict(data)[source]¶

Creates an ansatz from a dictionary. For deserialization purposes.

Parameters:: data (dict[str, ArrayLike])
Return type:: PolyExpAnsatz

concat(other, axis=0)[source]¶

Concatenates two PolyExpAnsatz objects along the specified batch axis.

All batch axes except the concatenation axis must agree in size. The lin_sup status must match. Core dimensions (num_CV_vars and num_derived_vars) must also match.

Parameters:

other (Ansatz) – The other PolyExpAnsatz to concatenate with.
axis (int) – The batch axis along which to concatenate (default 0).

Returns:

A new PolyExpAnsatz with the concatenated batch dimensions.

Raises:

TypeError – If other is not a PolyExpAnsatz.
ValueError – If batch dimensions don’t match (except at axis), if lin_sup status doesn’t match, or if core dimensions don’t match.

Return type:

PolyExpAnsatz

>>> from mrmustard.physics.ansatz import PolyExpAnsatz
>>> import numpy as np
>>> A1 = np.random.random((2, 3, 3))
>>> b1 = np.random.random((2, 3))
>>> c1 = np.random.random((2,))
>>> ansatz1 = PolyExpAnsatz(A1, b1, c1)
>>> A2 = np.random.random((3, 3, 3))
>>> b2 = np.random.random((3, 3))
>>> c2 = np.random.random((3,))
>>> ansatz2 = PolyExpAnsatz(A2, b2, c2)
>>> concatenated = ansatz1.concat(ansatz2, axis=0)
>>> assert concatenated.batch_shape == (5,)

contract(other, idxs)[source]¶

Contract along specified core (CV) axes, broadcasting batch dimensions and taking the kronecker product of linear-superposition dimensions.

Note

Batch dims broadcast automatically (no batch string needed).
Negative indices are normalized relative to each operand’s CV count.

Parameters:

other (Ansatz) – The other PolyExpAnsatz to contract with.
idxs (tuple[Sequence[int], Sequence[int]]) – Tuple (idx_self, idx_other) of CV-axis indices to integrate over (0-based relative to the CV variables). The two sequences must have the same length. Negative indices allowed.

Returns:

The contracted PolyExpAnsatz with kept core axes ordered as [self non-contracted] + [other non-contracted] and derived axes as [self derived] + [other derived].

Raises:

TypeError – If other is not a PolyExpAnsatz.

Return type:

PolyExpAnsatz

decompose_ansatz()[source]¶

This method decomposes a PolyExp ansatz to make it more efficient to evaluate. An ansatz with n CV variables and m derived variables has parameters with the following shapes: A=(batch;n+m,n+m), b=(batch;n+m), c = (batch;k_1,k_2,...,k_m;j_1,...,j_d), where d is the number of discrete variables, i.e. axes of the array of values that we get when we evaluate the ansatz at a point. This can be rewritten as an ansatz of dimension A=(*batch;2n,2n), b=(*batch;2n), c = (*batch;l_1,l_2,...,l_n;j_1,...,j_d), with l_i = sum_j k_j. This means that the number of continuous variables remains n, the number of derived variables decreases from m to n, and the number of discrete variables remains d. The price we pay is that the order of the derivatives is larger (the order of each derivative is the sum of all the orders of the initial derivatives). This decomposition is typically favourable if m > n and the sum of the elements in c.shape[1:] is not too large. This method will actually decompose the ansatz only if m > n and return the original ansatz otherwise.

Return type:: PolyExpAnsatz

eval(*z, batch_string=None)[source]¶

Evaluates the ansatz at given points or returns a partially evaluated ansatz. This method supports passing an einsum-style batch string to specify how the batch dimensions of the arguments should be handled. For example, for two arguments, “i,j->ij” means to take the outer product of the batch dimensions. The batch dimensions of the ansatz itself are not part of the batch string and are placed after the output batch dimensions in the output.

Partial evaluation: If any argument is None, it returns a new ansatz with those arguments unevaluated. For example, if F(z1, z2, z3) is called as F(1.0, None, 3.0), it returns a new ansatz G(z2) with z1 and z3 fixed at 1.0 and 3.0.

2. Full evaluation: If all arguments are provided, it returns the value of the ansatz at those points.

The returned shape depends on the batch shape of the ansatz itself and the batch dimensions of the inputs and the batch string.

Parameters:

z (Vector | None) – points in C where the function is (partially) evaluated or None if the variable is
evaluated. (not)
batch_string (str | None) – like einsum string for batch dimensions of the inputs, e.g. “i,j->ij”

Returns:

The value of the ansatz or a new ansatz if partial evaluation is performed.

Return type:

Scalar | ArrayLike | PolyExpAnsatz

reorder(order)[source]¶

Reorders the CV indices of an (A,b,c) triple. The length of order must equal the number of CV variables. This method returns a new ansatz object.

Parameters:: order (Sequence[int])
Return type:: PolyExpAnsatz

reorder_batch(order)[source]¶

Reorders the batch dimensions of the ansatz. The length of order must equal the number of batch dimensions. This method returns a new ansatz object.

Parameters:: order (Sequence[int]) – The desired order of the batch dimensions.
Returns:: A new Ansatz with reordered batch dimensions.
Return type:: PolyExpAnsatz

simplify()[source]¶

Returns a new ansatz simplified by combining terms that have the same exponential part, i.e. two components of the batch are considered equal if their A and b are equal. In this case only one is kept and the corresponding c are added.

Will return immediately if the ansatz has already been simplified, so it is safe to re-call.

Raises:: NotImplementedError – If the PolyExpAnsatz is batched (w/o linear superposition).
Return type:: PolyExpAnsatz

to_dict()[source]¶

Returns a dictionary representation of the ansatz. For serialization purposes.

Return type:: dict[str, ArrayLike]

trace(idx_z, idx_zconj, measure=-1.0)[source]¶

Computes the trace of the ansatz across the specified pairs of CV variables.

Parameters:

idx_z (tuple[int, ...]) – The indices indicating which CV variables to integrate over.
idx_zconj (tuple[int, ...]) – The indices indicating which conjugate CV variables to integrate over.
measure (float) – The measure to use in the complex Gaussian integral.

Returns:

A new ansatz with the specified indices traced out.

code/physics/ansatz

Download Python script

Download Notebook

View on GitHub

Ansatz¶

Inheritance Diagram¶

Ansatz¶

ArrayAnsatz¶

PolyExpAnsatz¶

Contents