Models

cymyc.approx.models

Model architectures for approximations of tensor fields on manifolds. The idea behind the form of each ansatz is that:

1. The tensor field is parameterised using a single (possibly vector-valued) function, \(\phi \in C^{\infty}(X)\).
2. The resulting tensor field should be globally defined over the manifold.

(2.) basically means that the local representation of the tensor field in each coordinate chart should glue together coherently. This particularly important for manifolds with a nontrivial topology - one appears to obtain nonsensical answers for downstream predictions if this is not respected.

The flax library is used here to define the models, but this is interchangable with any library or framework. As idiomatic jax is functional, the model definitions are kept separate from the parameters. During the forward pass, the parameters are passed separately as a dictionary to the apply method of the model.

There are multiple routes to add new architectures for approximation of various tensor fields. The simplest one, keeping in line with the logic in the models module, is to add:

1. A Flax module describing the sequence of operations defined by your architecture.

   import jax.numpy as jnp
   from flax import linen as nn
   
   class MyAnsatz(nn.Module):
   
   @nn.compact
   def __call__(self, local_coords):
       p = local_coords
       p_bar = jnp.conjugate(p)
       p_norm_sq = jnp.sum(p * p_bar)
       return jnp.outer(p, p_bar) / p_norm_sq
   

2. A pure function which accepts a pytree of parameters for the model and executes the computation by invoking the .apply method of the module you defined above.

   def ansatz_fun(p, params, *args):
       p = ...  # some logic
       model = MyAnsatz(*args)  # model constructor
       return model.apply({'params': params}, p)
   

LearnedVector_spectral_nn

Bases: Module

Spectral network implementation for hypersurfaces embedded in \(\mathbb{P}^n\). The model defined here is a simple feedforward network with a spectral embedding layer at the input, which converts points expressed in homogeneous coordinates to a \(\mathbb{C}^*\)-invariant matrix representation. The resulting parameterised function is globally defined over \(X\) - one may then construct various ansatze off this. For details see arxiv:2211.09801.

Attributes:

Name	Type	Description
dim	int	Dimension of projective space + 1.
ambient	Sequence[int]	Dimensions of the ambient space factors.
n_units	(Sequence[int], optional)	Number of units in each layer, by default (64, 64, 64).
n_out	(int, optional)	Number of output features, by default 1.
use_spectral_embedding	(bool, optional)	Whether to use spectral embedding, default True.
activation	(Callable[[ndarray], ndarray], optional)	Nonlinearity between each layer, default nn.gelu.
cy_dim	(int, optional)	Dimension of the complex projective space, default 3.

spectral_layer(x: Float[Array, i], x_dim: int) -> Array

Converts homogeneous \([z_0 : ... : z_n]\) coords in \(\mathbb{P}^n\) into the first-order basis of eigenfunctions of the Laplacian in projective space, whose generic form is \(z_i \bar{z}_j / \sum z_k \bar{z}_k\).

Parameters:

Name	Type	Description	Default
x	Float[Array, i]	Input array of homogeneous coordinates.	required
x_dim	int	Dimension of the input array.	required

Returns:

Type	Description
Complex[Array, i * i + 1 // 2]	First-order basis of eigenfunctions of the Laplacian as a vector.

__call__(x: Float[Array, i]) -> Array

Spectral NN forward pass for hypersurfaces embedded in a single projective space factor.

Parameters:

Name	Type	Description	Default
x	Float[Array, i]	Input array of homogeneous coordinates.	required

Returns:

Type	Description
Complex[Array, n_out]	Output of vector-valued function.

LearnedVector_spectral_nn_CICY

Bases: LearnedVector_spectral_nn

Spectral network implementation for manifolds embedded in a product of projective space factors, \({P}^{n_1} \times \cdots \times \mathbb{P}^{n_K}\).

__call__(x: Float[Array, i]) -> Array

Spectral NN forward pass for complete intersection manifolds in a product of projective spaces. The spectral layer converts the coordinates for each projective space to its respective \(\mathbb{C}^*\)-invariant matrix representation, and concatenates the results.

Parameters:

Name	Type	Description	Default
x	Float[Array, i]	Input array of homogeneous coordinates, local coordinates for each projective space factor are listed consecutively in the input array.	required

Returns:

Type	Description
Complex[Array, n_out]	Output of vector-valued function.

CoeffNetwork_spectral_nn_CICY

Bases: LearnedVector_spectral_nn

Spectral network parameterising the coefficients of a linear combination of a basis of sections for the holomorphic bundle \(V \rightarrow X\). The sections are globally defined by construction, hence the linear combination is a global section of \(V\), parameterised by a neural network. Schematically,

\[ s:= \sum_I \psi_I^{(\textsf{NN})} \cdot \mathbf{e}^I~. \]

Where summation over the multi-index \(I\) denotes contraction of appropriate tensor indices. Here the network specialises to the case of \(V=T_X\), the standard embedding, but interchangeable with any other bundle \(V \rightarrow X\) by subclassing and modifying the shapes of the coefficients.

The case of multiple parameterised sections \(s^{(k)} \in \Gamma(V)\) modelled is handled by the h_21 parameter - this sets the number of coefficients output by the network.

Inherits from LearnedVector_spectral_nn.

Attributes:

Name	Type	Description
h_21	int	Dimension of complex structure moduli space - controls the number of sections learnt.

__call__(x: Float[Array, i]) -> Array

Forward pass for coefficients, modelled as vector-valued functions on \(X\).

Parameters:

Name	Type	Description	Default
x	Float[Array, i]	Input array of homogeneous coordinates, local coordinates for each projective space factor are listed consecutively in the input array.	required

Returns:

Type	Description
Complex[Array, n_out]	Output of vector-valued function.

phi_head(p: Float[Array, i], params: Mapping[str, Array], n_hyper: int, ambient: Sequence[int], n_out: int = 1, spectral: bool = True, activation: Callable[[jnp.ndarray], jnp.ndarray] = nn.gelu) -> Complex[Array, n_out]

Wrapper to feed parameters into forward pass for \(\phi\)-component in the ddbar_phi_model.

Parameters:

Name	Type	Description	Default
p	Float[Array, i]	Input array of homogeneous coordinates.	required
params	Mapping[Str, Array]	Model parameters stored as a dictionary - keys are the module names registered upon initialisation and values are the parameter values.	required
n_hyper	int	Number of defining equations of the complete intersection.	required
ambient	Sequence[int]	Sequence representing the ambient space dimensions.	required
n_out	int	Dimension of output vector, default 1.	1
spectral	bool	Toggles spectral embedding, default True.	True
activation	Callable[[ndarray], ndarray]	Activation function, default nn.gelu.	gelu

Returns:

Type	Description
Complex[Array, n_out]	Output of vector-valued function.

ddbar_phi_model(p: Float[Array, i], params: Mapping[str, Array], g_ref_fn: Callable[[ArrayLike], jnp.ndarray], g_correction_fn: Callable[[ArrayLike], jnp.ndarray]) -> jnp.ndarray

Returns metric on \(X\) under pullback from ambient space as an \(\partial \bar{\partial}\)-exact correction to the reference Fubini-Study metric in the desired Kähler class, $$ \tilde{g} := g_{\text{ref}} + \partial \overline{\partial} \phi; \quad \phi \in C^{\infty}(X)~. $$ Generates pullbacks on the fly.

Parameters:

Name	Type	Description	Default
p	Float[Array, i]	Input array of homogeneous coordinates.	required
params	Mapping[Str, Array]	Model parameters stored as a dictionary - keys are the module names registered upon initialisation and values are the parameter values.	required
g_ref_fn	Callable[[ndarray], ndarray]	Function representing reference Fubini-Study metric in local coordinates.	required
g_correction_fn	Callable[[ndarray], ndarray]	Function \(\phi \in C^{\infty}(X)\) which generates the \(\partial \bar{\partial}\)-exact correction to the reference metric.	required

coeff_head(p: Float[Array, i], params: Mapping[str, Array], n_homo_coords: int, ambient: Sequence[int], h_21: int = 1, activation: Callable[[jnp.ndarray], jnp.ndarray] = nn.gelu) -> jnp.ndarray

Wrapper to feed parameters into forward pass for section coefficient network.

Parameters:

Name	Type	Description	Default
p	Float[Array, i]	Input array of homogeneous coordinates.	required
params	Mapping[Str, Array]	Model parameters stored as a dictionary - keys are the module names registered upon initialisation and values are the parameter values.	required
n_homo_coords	int	Number of homogeneous coordinates.	required
ambient	Sequence[int]	Sequence representing the ambient space dimensions.	required
h_21	int	Number of sections learnt / harmonic \((0,1)\) forms on \(X\). This controls the size of the coefficient network.	1
activation	Callable[[ndarray], ndarray]	Activation function, default nn.gelu.	gelu