IMLCV.implementations.CvDiscovery

Classes

Encoder	Base class for all neural network modules.
Decoder	Base class for all neural network modules.
VAE	Base class for all neural network modules.
TranformerAutoEncoder	Base class for dataclasses that should act like a JAX pytree node.
TransoformerLDA	Base class for dataclasses that should act like a JAX pytree node.
TransformerMAF	Base class for dataclasses that should act like a JAX pytree node.

Functions

_LDA_trans(cv, nl, shmap, shmap_kwargs, alpha, outdim, ...)
_LDA_rescale(cv, nl, shmap, shmap_kwargs, mean)
_scale_trans(cv, nl, shmap, shmap_kwargs, alpha, ...)

Module Contents

class IMLCV.implementations.CvDiscovery.Encoder

Bases: flax.linen.Module

Base class for all neural network modules.

Layers and models should subclass this class.

All Flax Modules are Python 3.7 dataclasses. Since dataclasses take over __init__, you should instead override setup(), which is automatically called to initialize the module.

Modules can contain submodules, and in this way can be nested in a tree structure. Submodels can be assigned as regular attributes inside the setup() method.

You can define arbitrary “forward pass” methods on your Module subclass. While no methods are special-cased, __call__ is a popular choice because it allows you to use module instances as if they are functions:

>>> from flax import linen as nn
>>> from typing import Tuple

>>> class Module(nn.Module):
...   features: Tuple[int, ...] = (16, 4)

...   def setup(self):
...     self.dense1 = nn.Dense(self.features[0])
...     self.dense2 = nn.Dense(self.features[1])

...   def __call__(self, x):
...     return self.dense2(nn.relu(self.dense1(x)))

Optionally, for more concise module implementations where submodules definitions are co-located with their usage, you can use the compact() wrapper.

latents: int

layers: int

nunits: int

dim: int

__call__(x)

class IMLCV.implementations.CvDiscovery.Decoder

Bases: flax.linen.Module

Base class for all neural network modules.

Layers and models should subclass this class.

All Flax Modules are Python 3.7 dataclasses. Since dataclasses take over __init__, you should instead override setup(), which is automatically called to initialize the module.

Modules can contain submodules, and in this way can be nested in a tree structure. Submodels can be assigned as regular attributes inside the setup() method.

You can define arbitrary “forward pass” methods on your Module subclass. While no methods are special-cased, __call__ is a popular choice because it allows you to use module instances as if they are functions:

>>> from flax import linen as nn
>>> from typing import Tuple

>>> class Module(nn.Module):
...   features: Tuple[int, ...] = (16, 4)

...   def setup(self):
...     self.dense1 = nn.Dense(self.features[0])
...     self.dense2 = nn.Dense(self.features[1])

...   def __call__(self, x):
...     return self.dense2(nn.relu(self.dense1(x)))

Optionally, for more concise module implementations where submodules definitions are co-located with their usage, you can use the compact() wrapper.

latents: int

layers: int

nunits: int

dim: int

__call__(z)

class IMLCV.implementations.CvDiscovery.VAE

Bases: flax.linen.Module

Base class for all neural network modules.

Layers and models should subclass this class.

All Flax Modules are Python 3.7 dataclasses. Since dataclasses take over __init__, you should instead override setup(), which is automatically called to initialize the module.

Modules can contain submodules, and in this way can be nested in a tree structure. Submodels can be assigned as regular attributes inside the setup() method.

You can define arbitrary “forward pass” methods on your Module subclass. While no methods are special-cased, __call__ is a popular choice because it allows you to use module instances as if they are functions:

>>> from flax import linen as nn
>>> from typing import Tuple

>>> class Module(nn.Module):
...   features: Tuple[int, ...] = (16, 4)

...   def setup(self):
...     self.dense1 = nn.Dense(self.features[0])
...     self.dense2 = nn.Dense(self.features[1])

...   def __call__(self, x):
...     return self.dense2(nn.relu(self.dense1(x)))

Optionally, for more concise module implementations where submodules definitions are co-located with their usage, you can use the compact() wrapper.

latents: int

layers: int

nunits: int

dim: int

setup()

Initializes a Module lazily (similar to a lazy __init__).

setup is called once lazily on a module instance when a module is bound, immediately before any other methods like __call__ are invoked, or before a setup-defined attribute on self is accessed.

This can happen in three cases:

1. Immediately when invoking apply(), init() or init_and_output().

2. Once the module is given a name by being assigned to an attribute of another module inside the other module’s setup method (see __setattr__()):

   >>> class MyModule(nn.Module):
   ...   def setup(self):
   ...     submodule = nn.Conv(...)
   
   ...     # Accessing `submodule` attributes does not yet work here.
   
   ...     # The following line invokes `self.__setattr__`, which gives
   ...     # `submodule` the name "conv1".
   ...     self.conv1 = submodule
   
   ...     # Accessing `submodule` attributes or methods is now safe and
   ...     # either causes setup() to be called once.
   

3. Once a module is constructed inside a method wrapped with compact(), immediately before another method is called or setup defined attribute is accessed.

__call__(x, z_rng)

encode(x)

classmethod reparameterize(rng, mean, logvar)

class IMLCV.implementations.CvDiscovery.TranformerAutoEncoder(*args, **kwargs)

Bases: IMLCV.base.CVDiscovery.Transformer

Base class for dataclasses that should act like a JAX pytree node.

nunits: int = 250

nlayers: int = 3

lr: float = 0.0001

num_epochs: int = 100

batch_size: int = 32

_fit(cv: list[IMLCV.base.CV.CV], cv_t: list[IMLCV.base.CV.CV], w: list[jax.Array], dlo: IMLCV.base.rounds.DataLoaderOutput, chunk_size=None, verbose=True, macro_chunk=1000)

IMLCV.implementations.CvDiscovery._LDA_trans(cv: IMLCV.base.CV.CV, nl: IMLCV.base.CV.NeighbourList | None, shmap, shmap_kwargs, alpha, outdim, solver)

IMLCV.implementations.CvDiscovery._LDA_rescale(cv: IMLCV.base.CV.CV, nl: IMLCV.base.CV.NeighbourList | None, shmap, shmap_kwargs, mean)

IMLCV.implementations.CvDiscovery._scale_trans(cv: IMLCV.base.CV.CV, nl: IMLCV.base.CV.NeighbourList | None, shmap, shmap_kwargs, alpha: jax.Array, scale_factor: jax.Array)

class IMLCV.implementations.CvDiscovery.TransoformerLDA(*args, **kwargs)

Bases: IMLCV.base.CVDiscovery.Transformer

Base class for dataclasses that should act like a JAX pytree node.

kernel = False

optimizer = None

solver: str = 'eigen'

method: str = 'pymanopt'

harmonic = True

min_gradient_norm: float = 0.001

min_step_size: float = 0.001

max_iterations: int = 25

_fit(cv_list: list[IMLCV.base.CV.CV], cv_t_list: list[IMLCV.base.CV.CV], w: list[jax.Array], dlo: IMLCV.base.rounds.DataLoaderOutput, chunk_size=None, verbose=True, macro_chunk=1000)

class IMLCV.implementations.CvDiscovery.TransformerMAF(*args, **kwargs)

Bases: IMLCV.base.CVDiscovery.Transformer

Base class for dataclasses that should act like a JAX pytree node.

eps: float = 1e-05

eps_pre: float = 1e-05

max_features: int = 500

max_features_pre: int = 500

sym: bool = True

use_w: bool = True

min_t_frac: float = 0.1

trans: IMLCV.base.CV.CvTrans | None = None

T_scale: float = 1.0

_fit(x: list[IMLCV.base.CV.CV] | list[IMLCV.base.CV.SystemParams], x_t: list[IMLCV.base.CV.CV] | list[IMLCV.base.CV.SystemParams] | None, w: list[jax.Array], w_t: list[jax.Array], dlo: IMLCV.base.rounds.DataLoaderOutput, macro_chunk=1000, chunk_size=None, verbose=True)