jinns 1.2.0__py3-none-any.whl → 1.3.0__py3-none-any.whl

jinns/utils/_spinn.py

@@ -1,249 +0,0 @@
-"""
-Implements utility function to create Separable PINNs
-https://arxiv.org/abs/2211.08761
-"""
-
-from dataclasses import InitVar
-from typing import Callable, Literal
-import jax
-import jax.numpy as jnp
-import equinox as eqx
-from jaxtyping import Key, Array, Float, PyTree
-
-from jinns.parameters._params import Params, ParamsDict
-
-
-class _SPINN(eqx.Module):
-    """
-    Construct a Separable PINN as proposed in
-    Cho et al., _Separable Physics-Informed Neural Networks_, NeurIPS, 2023
-
-    Parameters
-    ----------
-    key : InitVar[Key]
-        A jax random key for the layer initializations.
-    d : int
-        The number of dimensions to treat separately, including time `t` if
-        used for non-stationnary equations.
-    eqx_list : InitVar[tuple[tuple[Callable, int, int] | Callable, ...]]
-        A tuple of tuples of successive equinox modules and activation functions to
-        describe the PINN architecture. The inner tuples must have the eqx module or
-        activation function as first item, other items represents arguments
-        that could be required (eg. the size of the layer).
-        The `key` argument need not be given.
-        Thus typical example is `eqx_list=
-        ((eqx.nn.Linear, 2, 20),
-            jax.nn.tanh,
-            (eqx.nn.Linear, 20, 20),
-            jax.nn.tanh,
-            (eqx.nn.Linear, 20, 20),
-            jax.nn.tanh,
-            (eqx.nn.Linear, 20, 1)
-        )`.
-    """
-
-    d: int = eqx.field(static=True, kw_only=True)
-
-    key: InitVar[Key] = eqx.field(kw_only=True)
-    eqx_list: InitVar[tuple[tuple[Callable, int, int] | Callable, ...]] = eqx.field(
-        kw_only=True
-    )
-
-    layers: list = eqx.field(init=False)
-    separated_mlp: list = eqx.field(init=False)
-
-    def __post_init__(self, key, eqx_list):
-        self.separated_mlp = []
-        for _ in range(self.d):
-            self.layers = []
-            for l in eqx_list:
-                if len(l) == 1:
-                    self.layers.append(l[0])
-                else:
-                    key, subkey = jax.random.split(key, 2)
-                    self.layers.append(l[0](*l[1:], key=subkey))
-            self.separated_mlp.append(self.layers)
-
-    def __call__(
-        self, inputs: Float[Array, "dim"] | Float[Array, "dim+1"]
-    ) -> Float[Array, "d embed_dim*output_dim"]:
-        outputs = []
-        for d in range(self.d):
-            t_ = inputs[d : d + 1]
-            for layer in self.separated_mlp[d]:
-                t_ = layer(t_)
-            outputs += [t_]
-        return jnp.asarray(outputs)
-
-
-class SPINN(eqx.Module):
-    """
-    A SPINN object compatible with the rest of jinns.
-    This is typically created with `create_SPINN`.
-
-    Parameters
-    ----------
-    d : int
-        The number of dimensions to treat separately, including time `t` if
-        used for non-stationnary equations.
-
-    """
-
-    d: int = eqx.field(static=True, kw_only=True)
-    r: int = eqx.field(static=True, kw_only=True)
-    eq_type: str = eqx.field(static=True, kw_only=True)
-    m: int = eqx.field(static=True, kw_only=True)
-
-    spinn_mlp: InitVar[eqx.Module] = eqx.field(kw_only=True)
-
-    params: PyTree = eqx.field(init=False)
-    static: PyTree = eqx.field(init=False, static=True)
-
-    def __post_init__(self, spinn_mlp):
-        self.params, self.static = eqx.partition(spinn_mlp, eqx.is_inexact_array)
-
-    @property
-    def init_params(self) -> PyTree:
-        """
-        Returns an initial set of parameters
-        """
-        return self.params
-
-    def __call__(
-        self,
-        t_x: Float[Array, "batch_size 1+dim"],
-        params: Params | ParamsDict | PyTree,
-    ) -> Float[Array, "output_dim"]:
-        """
-        Evaluate the SPINN on some inputs with some params.
-        """
-        try:
-            spinn = eqx.combine(params.nn_params, self.static)
-        except (KeyError, AttributeError, TypeError) as e:
-            spinn = eqx.combine(params, self.static)
-        v_model = jax.vmap(spinn)
-        res = v_model(t_x)
-
-        a = ", ".join([f"{chr(97 + d)}z" for d in range(res.shape[1])])
-        b = "".join([f"{chr(97 + d)}" for d in range(res.shape[1])])
-        res = jnp.stack(
-            [
-                jnp.einsum(
-                    f"{a} -> {b}",
-                    *(
-                        res[:, d, m * self.r : (m + 1) * self.r]
-                        for d in range(res.shape[1])
-                    ),
-                )
-                for m in range(self.m)
-            ],
-            axis=-1,
-        )  # compute each output dimension
-
-        # force (1,) output for non vectorial solution (consistency)
-        if len(res.shape) == self.d:
-            return jnp.expand_dims(res, axis=-1)
-        return res
-
-
-def create_SPINN(
-    key: Key,
-    d: int,
-    r: int,
-    eqx_list: tuple[tuple[Callable, int, int] | Callable, ...],
-    eq_type: Literal["ODE", "statio_PDE", "nonstatio_PDE"],
-    m: int = 1,
-) -> tuple[SPINN, PyTree]:
-    """
-    Utility function to create a SPINN neural network with the equinox
-    library.
-
-    *Note* that a SPINN is not vmapped and expects the
-    same batch size for each of its input axis. It directly outputs a solution
-    of shape `(batchsize, batchsize)`. See the paper for more details.
-
-    Parameters
-    ----------
-    key
-        A JAX random key that will be used to initialize the network parameters
-    d
-        The number of dimensions to treat separately.
-    r
-        An integer. The dimension of the embedding.
-    eqx_list
-        A tuple of tuples of successive equinox modules and activation functions to
-        describe the PINN architecture. The inner tuples must have the eqx module or
-        activation function as first item, other items represents arguments
-        that could be required (eg. the size of the layer).
-        The `key` argument need not be given.
-        Thus typical example is
-        `eqx_list=((eqx.nn.Linear, 2, 20),
-            jax.nn.tanh,
-            (eqx.nn.Linear, 20, 20),
-            jax.nn.tanh,
-            (eqx.nn.Linear, 20, 20),
-            jax.nn.tanh,
-            (eqx.nn.Linear, 20, 1)
-        )`.
-    eq_type : Literal["ODE", "statio_PDE", "nonstatio_PDE"]
-        A string with three possibilities.
-        "ODE": the PINN is called with one input `t`.
-        "statio_PDE": the PINN is called with one input `x`, `x`
-        can be high dimensional.
-        "nonstatio_PDE": the PINN is called with two inputs `t` and `x`, `x`
-        can be high dimensional.
-        **Note**: the input dimension as given in eqx_list has to match the sum
-        of the dimension of `t` + the dimension of `x` or the output dimension
-        after the `input_transform` function.
-    m
-        The output dimension of the neural network. According to
-        the SPINN article, a total embedding dimension of `r*m` is defined. We
-        then sum groups of `r` embedding dimensions to compute each output.
-        Default is 1.
-
-
-
-    Returns
-    -------
-    spinn
-        An instanciated SPINN
-    spinn.init_params
-        The initial set of parameters of the model
-
-    Raises
-    ------
-    RuntimeError
-        If the parameter value for eq_type is not in `["ODE", "statio_PDE",
-        "nonstatio_PDE"]` and for various failing checks
-    """
-
-    if eq_type not in ["ODE", "statio_PDE", "nonstatio_PDE"]:
-        raise RuntimeError("Wrong parameter value for eq_type")
-
-    try:
-        nb_inputs_declared = eqx_list[0][1]  # normally we look for 2nd ele of 1st layer
-    except IndexError:
-        nb_inputs_declared = eqx_list[1][
-            1
-        ]  # but we can have, eg, a flatten first layer
-    if nb_inputs_declared != 1:
-        raise ValueError("Input dim must be set to 1 in SPINN!")
-
-    try:
-        nb_outputs_declared = eqx_list[-1][2]  # normally we look for 3rd ele of
-        # last layer
-    except IndexError:
-        nb_outputs_declared = eqx_list[-2][2]
-        # but we can have, eg, a `jnp.exp` last layer
-    if nb_outputs_declared != r * m:
-        raise ValueError("Output dim must be set to r * m in SPINN!")
-
-    if d > 24:
-        raise ValueError(
-            "Too many dimensions, not enough letters available in jnp.einsum"
-        )
-
-    spinn_mlp = _SPINN(key=key, d=d, eqx_list=eqx_list)
-    spinn = SPINN(spinn_mlp=spinn_mlp, d=d, r=r, eq_type=eq_type, m=m)
-
-    return spinn, spinn.init_params

jinns-1.2.0.dist-info/RECORD

@@ -1,41 +0,0 @@
-jinns/__init__.py,sha256=5p7V5VJd7PXEINqhqS4mUsnQtXlyPwfctRhL4p0loFg,181
-jinns/data/_Batchs.py,sha256=oc7-N1wEbsEvbe9fjVFKG2OPoZJVEjzPm8uj_icACf4,817
-jinns/data/_DataGenerators.py,sha256=BWNF28jM7EN6wyJ3X2X1qjG4sjs6vHLgIVVY6aqhKS8,65099
-jinns/data/__init__.py,sha256=TRCH0Z4-SQZ50MbSf46CUYWBkWVDmXCyez9T-EGiv_8,338
-jinns/experimental/__init__.py,sha256=3jCIy2R2i_0Erwxg-HwISdH79Nt1XCXhS9yY1F5awiY,208
-jinns/experimental/_diffrax_solver.py,sha256=upMr3kTTNrxEiSUO_oLvCXcjS9lPxSjvbB81h3qlhaU,6813
-jinns/loss/_DynamicLoss.py,sha256=x1pJLZ_4P2iHhf5hjRhkWi1v5Q2dAH-v7Gctv4ax73E,25819
-jinns/loss/_DynamicLossAbstract.py,sha256=bqmPxyrcvZh_dL74DTpj-TGiFxchvG8qC6KhuGeyOoA,12006
-jinns/loss/_LossODE.py,sha256=lTJ3b2EnlYk1eAcldiVyFRh_XQrk83eYRQBuqWSiILg,23418
-jinns/loss/_LossPDE.py,sha256=ufECwKdQXFEy8h5fdDA2uFTjiO_sVdqi9nfP7NTyu60,48775
-jinns/loss/__init__.py,sha256=PRiJV9fd2GSwaCBVCPyh6pFc6pdA40jfb_T1YvO8ERc,712
-jinns/loss/_boundary_conditions.py,sha256=KL_UUajQWLtRFqmBO_lfEcxrUsY-398ySTgpfaHdeYk,12277
-jinns/loss/_loss_utils.py,sha256=RXu20VYeAelpuRIAJ67k3h7jdLsu2VIdjzmLVdiMjns,14223
-jinns/loss/_loss_weights.py,sha256=F0Fgji2XpVk3pr9oIryGuXcG1FGQo4Dv6WFgze2BtA0,2201
-jinns/loss/_operators.py,sha256=Hf5zWUQCGRSm1SS8f6IZLqC8bL7suC-R4ckXNbXqoAI,21790
-jinns/parameters/__init__.py,sha256=1gxNLoAXUjhUzBWuh86YjU5pYy8SOboCs8TrKcU1wZc,158
-jinns/parameters/_derivative_keys.py,sha256=UyEcgfNF1vwPcGWD2ShAZkZiq4thzRDm_OUJzOfjjiY,21909
-jinns/parameters/_params.py,sha256=wK9ZSqoL9KnjOWqc_ZhJ09ffbsgeUEcttc1Rhme0lLk,3550
-jinns/plot/__init__.py,sha256=Q279h5veYWNLQyttsC8_tDOToqUHh8WaRON90CiWXqk,81
-jinns/plot/_plot.py,sha256=jS0fdI9VHCd6HL4K8ZzDjmMr10EJFqU6a0S4iWpEzFk,11433
-jinns/solver/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-jinns/solver/_rar.py,sha256=BCvUdNLicdU1YizHKcSaJLmkUkCybyGmTRjHnjoiQKA,9730
-jinns/solver/_solve.py,sha256=TR-RuH2tYIFq6YhbZ-Aa7r8OYy7CstjEYNm2RMI3xck,22641
-jinns/solver/_utils.py,sha256=b2zYvwZY_fU0NMNWvUEMvHez9s7hwcxfpGzQlz5F6HA,5762
-jinns/utils/__init__.py,sha256=nx7wi4RfgeHr5wLM_c1BAG8UQteTUTK6B0t9enmTn1E,243
-jinns/utils/_containers.py,sha256=a7A-iUApnjc1YVc7bdt9tKUvHHPDOKMB9OfdrDZGWN8,1450
-jinns/utils/_hyperpinn.py,sha256=hcdHRHXer03I3Twupok5moPR1vRGgPQUCjFHAZMCMdc,16872
-jinns/utils/_pinn.py,sha256=8wDliUrDbBUKWllJgllapNX5BLLHMWF7Mz2CaEA-6vk,12667
-jinns/utils/_ppinn.py,sha256=K4xuGaT-4XddIxsi1wOYO24G9l1RAY2m4bZqP7C3klI,8690
-jinns/utils/_save_load.py,sha256=BngwUjkm5jibctnUCB8klMypYkfE6DKLaoWehwWnTNU,8585
-jinns/utils/_spinn.py,sha256=tjly3qNZj9C2JntFKuwrLfFAbVfvy7Z5HPjj0ymfnLk,8337
-jinns/utils/_types.py,sha256=EQXxUAipwy7lVk39hYw4NDf_a66OP01w_T3ziPITxpU,1901
-jinns/utils/_utils.py,sha256=hoRcJqcTuQi_Ip40oI4EbxW46E1rp2C01_HfuCpwKRM,2932
-jinns/validation/__init__.py,sha256=Jv58mzgC3F7cRfXA6caicL1t_U0UAhbwLrmMNVg6E7s,66
-jinns/validation/_validation.py,sha256=bvqL2poTFJfn9lspWqMqXvQGcQIodKwKrC786QtEZ7A,4700
-jinns-1.2.0.dist-info/AUTHORS,sha256=7NwCj9nU-HNG1asvy4qhQ2w7oZHrn-Lk5_wK_Ve7a3M,80
-jinns-1.2.0.dist-info/LICENSE,sha256=BIAkGtXB59Q_BG8f6_OqtQ1BHPv60ggE9mpXJYz2dRM,11337
-jinns-1.2.0.dist-info/METADATA,sha256=PxuzXslHG4YzOcuwDY3ulo9zJaF95ChjK9y584eMnN0,4662
-jinns-1.2.0.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
-jinns-1.2.0.dist-info/top_level.txt,sha256=RXbkr2hzy8WBE8aiRyrJYFqn3JeMJIhMdybLjjLTB9c,6
-jinns-1.2.0.dist-info/RECORD,,