jinnax.nn

   1#Functions to train NN
   2import jax
   3import jax.numpy as jnp
   4import optax
   5from alive_progress import alive_bar
   6import math
   7import time
   8import pandas as pd
   9import numpy as np
  10import matplotlib.pyplot as plt
  11import dill as pickle
  12from genree import bolstering as gb
  13from genree import kernel as gk
  14from jax import random
  15from jinnax import data as jd
  16import os
  17import numpy as np
  18from scipy.fft import dst, idst
  19from itertools import product
  20from functools import partial
  21import orthax
  22from jax import lax
  23from jaxopt import LBFGS
  24
  25__docformat__ = "numpy"
  26
  27#MSE
  28@jax.jit
  29def MSE(pred,true):
  30    """
  31    Squared error
  32    ----------
  33    Parameters
  34    ----------
  35    pred : jax.numpy.array
  36
  37        A JAX numpy array with the predicted values
  38
  39    true : jax.numpy.array
  40
  41        A JAX numpy array with the true values
  42
  43    Returns
  44    -------
  45    squared error
  46    """
  47    return (true - pred) ** 2
  48
  49#MSE self-adaptative
  50@jax.jit
  51def MSE_SA(pred,true,w,q = 2):
  52    """
  53    Self-adaptative squared error
  54    ----------
  55    Parameters
  56    ----------
  57    pred : jax.numpy.array
  58
  59        A JAX numpy array with the predicted values
  60
  61    true : jax.numpy.array
  62
  63        A JAX numpy array with the true values
  64
  65    weight : jax.numpy.array
  66
  67        A JAX numpy array with the weights
  68
  69    q : float
  70
  71        Power for the weights mask
  72
  73    Returns
  74    -------
  75    self-adaptative squared error with polynomial mask
  76    """
  77    return (w ** q) * ((true - pred) ** 2)
  78
  79#L2 error
  80@jax.jit
  81def L2error(pred,true):
  82    """
  83    L2-error in percentage (%)
  84    ----------
  85    Parameters
  86    ----------
  87    pred : jax.numpy.array
  88
  89        A JAX numpy array with the predicted values
  90
  91    true : jax.numpy.array
  92
  93        A JAX numpy array with the true values
  94
  95    Returns
  96    -------
  97    L2-error
  98    """
  99    return 100*jnp.sqrt(jnp.sum((true - pred)**2))/jnp.sqrt(jnp.sum(true ** 2))
 100
 101#Auxialiry functions to sample singular Matern
 102def idst1(x,axis = -1):
 103    """
 104    Inverse Discrete Sine Transform of type I with orthonormal scaling
 105    ----------
 106    Parameters
 107    ----------
 108    x : jax.numpy.array
 109
 110        Array to apply the transformation
 111
 112    axis : int
 113
 114        Axis to apply the transformation over
 115
 116    Returns
 117    -------
 118    jax.numpy.array
 119    """
 120    return idst(x,type = 1,axis = axis,norm = 'ortho')
 121
 122def dstn(x,axes = None):
 123    """
 124    Discrete Sine Transform of type I with orthonormal scaling over many axes
 125    ----------
 126    Parameters
 127    ----------
 128    x : jax.numpy.array
 129
 130        Array to apply the transformation
 131
 132    axes : int
 133
 134        Axes to apply the transformation over
 135
 136    Returns
 137    -------
 138    jax.numpy.array
 139    """
 140    if axes is None:
 141        axes = tuple(range(x.ndim))
 142    y = x
 143    for ax in axes:
 144        y = dst(x,type = 1,axis = ax,norm = 'ortho')
 145    return y
 146
 147def idstn(x,axes = None):
 148    """
 149    Inverse Discrete Sine Transform of type I with orthonormal scaling over many axes
 150    ----------
 151    Parameters
 152    ----------
 153    x : jax.numpy.array
 154
 155        Array to apply the transformation
 156
 157    axes : int
 158
 159        Axes to apply the transformation over
 160
 161    Returns
 162    -------
 163    jax.numpy.array
 164    """
 165    if axes is None:
 166        axes = tuple(range(x.ndim))
 167    y = x
 168    for ax in axes:
 169        y = idst1(y,axis = ax)
 170    return y
 171
 172def dirichlet_eigs_nd(n,L):
 173    """
 174    Eigenvalues of the discrete Dirichlet-Laplace operator in a rectangle
 175    ----------
 176    Parameters
 177    ----------
 178    n : list
 179
 180        List with the number of points in the grid in each dimension
 181
 182    L : list
 183
 184        List with the upper limit of the interval of the domain in each dimension. Assumed the lower limit is zero
 185
 186    Returns
 187    -------
 188    jax.numpy.array
 189    """
 190    #Unidimensional eigenvalues
 191    lam_axes = []
 192    for ni, Li in zip(n,L):
 193        h = Li / (ni + 1.0)
 194        k = jnp.arange(1,ni + 1,dtype = np.float32)
 195        ln = (2.0 / (h*h)) * (1.0 - jnp.cos(jnp.pi * k / (ni + 1.0)))
 196        lam_axes.append(ln)
 197    grids = jnp.meshgrid(*lam_axes, indexing='ij')
 198    Lam = jnp.zeros_like(grids[0])
 199    for g in grids:
 200        Lam += g
 201    return Lam
 202
 203
 204#Sample from d-dimensional Matern process
 205def generate_matern_sample(key,d = 2,N = 128,L = 1.0,kappa = 1,alpha = 1,sigma = 1,periodic = False):
 206    """
 207    Sample d-dimensional Matern process
 208    ----------
 209    Parameters
 210    ----------
 211    key : int
 212
 213        Seed for randomization
 214
 215    d : int
 216
 217        Dimension. Default 2
 218
 219    N : int
 220
 221        Size of grid in each dimension. Default 128
 222
 223    L : list of float
 224
 225        The domain of the function in each coordinate is [0,L[1]]. If a float, repeat the same interval for all coordinates. Default 1
 226
 227    kappa,alpha,sigma : float
 228
 229        Parameters of the Matern process
 230
 231    periodic : logical
 232
 233        Whether to sample with periodic boundary conditions. Periodic = False is not JAX native and does not work with JIT
 234
 235    Returns
 236    -------
 237    jax.numpy.array
 238    """
 239    if periodic:
 240        #Shape and key
 241        key = jax.random.PRNGKey(key)
 242        shape = (N,) * d
 243        if isinstance(L,float) or isinstance(L,int):
 244            L = d*[L]
 245        if isinstance(N,float) or isinstance(N,int):
 246            N = d*[N]
 247
 248        #Setup Frequency Grid (2D)
 249        freq = [jnp.fft.fftfreq(N[j],d = L[j]/N[j]) * 2 * jnp.pi for j in range(d)]
 250        grids = jnp.meshgrid(*freq, indexing='ij')
 251        sq_norm_xi = sum(g**2 for g in grids)
 252
 253        #Generate White Noise in Fourier Space
 254        key_re, key_im = jax.random.split(key)
 255        white_noise_f = (jax.random.normal(key_re, shape) +
 256                         1j * jax.random.normal(key_im, shape))
 257
 258        #Apply the Whittle Filter
 259        amplitude_filter = (kappa ** 2 + sq_norm_xi) ** (-alpha / 2.0)
 260        field_f = white_noise_f * amplitude_filter
 261
 262        #Transform back to Physical Space
 263        sample = jnp.real(jnp.fft.ifftn(field_f))
 264        return sigma*sample
 265    else: #NOT JAX
 266        #Shape and key
 267        rng = np.random.default_rng(seed = key)
 268        if isinstance(L,float) or isinstance(L,int):
 269            L = d*[L]
 270        if isinstance(N,float) or isinstance(N,int):
 271            N = d*[N]
 272        shape = tuple(N)
 273
 274        #White noise in real space
 275        W = rng.standard_normal(size = shape)
 276
 277        #To Dirichlet eigenbasis via separable DST-I (orthonormal)
 278        W_hat = dstn(W)
 279
 280        #Discrete Dirichlet Laplacian eigenvalues
 281        lam = dirichlet_eigs_nd(N, L)
 282
 283        #Spectral filter
 284        filt = ((kappa + lam) ** (-alpha/2.0))
 285        psi_hat = filt * W_hat
 286
 287        #Back to real space
 288        psi = idstn(psi_hat)
 289        return jnp.array(sigma*psi)
 290
 291#Vectorized generate_matern_sample
 292def generate_matern_sample_batch(d = 2,N = 512,L = 1.0,kappa = 10.0,alpha = 1,sigma = 10,periodic = False):
 293    """
 294    Create function to sample d-dimensional Matern process
 295    ----------
 296    Parameters
 297    ----------
 298    d : int
 299
 300        Dimension. Default 2
 301
 302    N : int
 303
 304        Size of grid in each dimension. Default 128
 305
 306    L : list of float
 307
 308        The domain of the function in each coordinate is [0,L[1]]. If a float, repeat the same interval for all coordinates. Default 1
 309
 310    kappa,alpha,sigma : float
 311
 312        Parameters of the Matern process
 313
 314    periodic : logical
 315
 316        Whether to sample with periodic boundary conditions. Periodic = False is not JAX native and does not work with JIT
 317
 318    Returns
 319    -------
 320    function
 321    """
 322    if periodic:
 323        return jax.vmap(lambda k: generate_matern_sample(k,d = d,N = N,L = L,kappa = kappa,alpha = alpha,sigma = sigma,periodic = periodic))
 324    else:
 325        return lambda keys: jnp.array(np.apply_along_axis(lambda k: generate_matern_sample(k,d = d,N = N,L = L,kappa = kappa,alpha = alpha,sigma = sigma,periodic = periodic),1,keys.reshape((keys.shape[0],1))))
 326
 327#Build function to compute the eigenfunctions of Laplacian
 328def eigenf_laplace(L_vec,kmax_per_axis = None,bc = "dirichlet",max_ef = None):
 329    """
 330    Create function to compute in batches the eigenfunctions of the Dirichlet-Laplace or Neumann-Laplace.
 331    ----------
 332    Parameters
 333    ----------
 334    L_vec : list of float
 335
 336        The domain of the function in each coordinate is [0,L[1]]
 337
 338    kmax_per_axis : list
 339
 340        List with the maximum number of eigenfunctions per dimension. Consider d * max(kmax_per_axis) eigenfunctions
 341
 342    bc : str
 343
 344        Boundary condition. 'dirichlet' or 'neumann'
 345
 346    max_ef : int
 347
 348        Maximum number of eigenfunctions to consider among the ones with greatest eigenvalues. If None, considers d * max(kmax_per_axis) eigenfunctions
 349
 350    Returns
 351    -------
 352    function to compute eigenfunctions,eigenvalues of the eigenfunctions considered
 353    """
 354    #Parameters
 355    L_vec = jnp.asarray(L_vec,dtype = jnp.float32)
 356    d = L_vec.shape[0]
 357    bc = bc.lower()
 358
 359    #Maximum number of functions
 360    if max_ef is None:
 361        if d == 1:
 362            max_ef = jnp.max(jnp.array(kmax_per_axis))
 363        else:
 364            max_ef = jnp.max(d * jnp.array(kmax_per_axis))
 365
 366    #Build the candidate multi-indices per axis
 367    kmax_per_axis = list(map(int, kmax_per_axis))
 368    if bc.startswith("d"):
 369        axis_ranges = [range(1, km + 1) for km in kmax_per_axis]
 370    elif bc.startswith("n"):
 371        axis_ranges = [range(0, km + 1) for km in kmax_per_axis]
 372
 373    #Get all multi-indices
 374    Ks_list = list(product(*axis_ranges))
 375    Ks = jnp.array(Ks_list,dtype = jnp.float32)
 376
 377    #Eigenvalues of the continuous Laplacian
 378    pi_over_L = jnp.pi / L_vec
 379    lambdas_all = jnp.sum((Ks * pi_over_L) ** 2, axis=1)
 380
 381    #Sort by eigenvalue
 382    order = jnp.argsort(lambdas_all)
 383    Ks = Ks[order]
 384    lambdas_all = lambdas_all[order]
 385
 386    #Keep first max_ef
 387    Ks = Ks[:max_ef]
 388    lambdas = lambdas_all[:max_ef]
 389    m = Ks.shape[0]
 390
 391    #Precompute per-feature normalization factor (closed form)
 392    def per_axis_norm_factor(k_i, L_i, is_dirichlet):
 393        if is_dirichlet:
 394            return jnp.sqrt(2.0 / L_i)
 395        else:
 396            return jnp.where(k_i == 0, jnp.sqrt(1.0 / L_i), jnp.sqrt(2.0 / L_i))
 397    if bc.startswith("d"):
 398        nf = jnp.prod(jnp.sqrt(2.0 / L_vec)[None, :],axis = 1)
 399        norm_factors = jnp.ones((m,),dtype = jnp.float32) * nf
 400    else:
 401        # per-mode product across axes
 402        def nf_row(k_row):
 403            return jnp.prod(per_axis_norm_factor(k_row, L_vec, False))
 404        norm_factors = jax.vmap(nf_row)(Ks)
 405
 406    #Build the callable function
 407    Ks_int = Ks  # float array, but only integer values
 408    L_vec_f = L_vec
 409    @jax.jit
 410    def phi(x):
 411        x = jnp.asarray(x,dtype = jnp.float32)
 412        #Initialize with ones
 413        vals = jnp.ones(x.shape[:-1] + (m,), dtype=jnp.float32)
 414        #Compute eigenfunction
 415        for i in range(d):
 416            ang = (jnp.pi / L_vec_f[i]) * x[..., i][..., None] * Ks_int[:, i]
 417            if bc.startswith("d"):
 418                comp = jnp.sin(ang)
 419            else:
 420                comp = jnp.cos(ang)
 421            vals = vals * comp
 422        #Apply L2-normalizing constants
 423        vals = vals * norm_factors[None, ...] if vals.ndim > 1 else vals * norm_factors
 424        return vals
 425    return phi, lambdas
 426
 427#Compute multiple frequences of domain aware fourrier fesatures
 428def multiple_daff(L_vec,kmax_per_axis = None,bc = "dirichlet",max_ef = None):
 429    """
 430    Create function to compute multiple frequences of the eigenfunctions of the Dirichlet-Laplace or Neumann-Laplace. Each frequences is a different domain.
 431    ----------
 432    Parameters
 433    ----------
 434    L_vec : list of lists of float
 435
 436        List with the domain of each frequence of the eigenfunctions in the form [0,L[i][1]]
 437
 438    kmax_per_axis : list
 439
 440        List with the maximum number of eigenfunctions per dimension.
 441
 442    bc : str
 443
 444        Boundary condition. 'dirichlet' or 'neumann'
 445
 446    max_ef : int
 447
 448        Maximum number of eigenfunctions to consider among the ones with greatest eigenvalues. If None, considers d * max(kmax_per_axis) eigenfunctions
 449
 450    Returns
 451    -------
 452    function to compute daff,eigenvalues of the eigenfunctions considered
 453    """
 454    psi = []
 455    lamb = []
 456    for L in L_vec:
 457        tmp,l = eigenf_laplace(L,kmax_per_axis,bc,max_ef) #Get function
 458        lamb.append(l)
 459        psi.append(tmp)
 460        del tmp
 461    #Create function to compute features
 462    @jax.jit
 463    def mff(x):
 464        y = []
 465        for i in range(len(psi)):
 466            y.append(psi[i](x))
 467        if len(psi) == 1:
 468            return y[0]
 469        else:
 470            return jnp.concatenate(y,1)
 471    return mff,jnp.concatenate(lamb)
 472
 473#Code for chebyshev polynomials writeen by AI (deprecated)
 474def _chebyshev_T_all(t, K: int):
 475    """
 476    Compute T_0..T_K(t) with the standard recurrence.
 477    t shape should be (..., d). We DO NOT squeeze any axis to preserve 'd'
 478    even when d == 1.
 479    Returns: array of shape (K+1, ...) matching t's batch dims, including d.
 480    """
 481    # Expect t to have last axis = d (keep it, even if d == 1)
 482    T0 = jnp.ones_like(t)           # (..., d)
 483    if K == 0:
 484        return T0[None, ...]        # (1, ..., d)
 485
 486    T1 = t                          # (..., d)
 487    if K == 1:
 488        return jnp.stack([T0, T1], axis=0)  # (2, ..., d)
 489
 490    def body(carry, _):
 491        Tkm1, Tk = carry            # each (..., d)
 492        Tkp1 = 2.0 * t * Tk - Tkm1  # (..., d)
 493        return (Tk, Tkp1), Tkp1
 494
 495    # K >= 2: produce T_2..T_K
 496    (_, _), T2_to_TK = lax.scan(body, (T0, T1), jnp.arange(K - 1))  # (K-1, ..., d)
 497    return jnp.concatenate([T0[None, ...], T1[None, ...], T2_to_TK], axis=0)  # (K+1, ..., d)
 498
 499@partial(jax.jit,static_argnums=(2,))  # n is static here; compile once per n
 500def multiple_cheb_fast(x, L_vec, n: int):
 501    """
 502    x: (N, d)
 503    L_vec: (L, d) containing 'b' endpoints (a is 0) for each dimension
 504    n: number of k terms (static)
 505    returns: (N, L*n)
 506    """
 507    N, d = x.shape
 508    L = L_vec.shape[0]
 509
 510    a = 0.0
 511    b = L_vec                       # (L, d)
 512    # Map x to t in [-1, 1] for each l, j: shape (L, N, d)
 513    t = (2.0 * x[None, :, :] - (a + b)[:, None, :]) / (b - a)[:, None, :]
 514
 515    # Chebyshev T_0..T_{n+2} for all (L, N, d): shape (n+3, L, N, d)
 516    T = _chebyshev_T_all(t, n + 2)
 517
 518    # phi_k = T_{k+2} - T_k, k = 0..n-1  => shape (n, L, N, d)
 519    ks = jnp.arange(n)
 520    phi = T[ks + 2, ...] - T[ks, ...]
 521
 522    # Multiply across dimensions (over the last axis = d) => (n, L, N)
 523    z = jnp.prod(phi, axis=-1)
 524
 525    # Reorder to (N, L, n) then flatten to (N, L*n)
 526    z = jnp.transpose(z, (2, 1, 0)).reshape(N, L * n)
 527    return z
 528
 529def multiple_cheb(L_vec, n: int):
 530    """
 531    Factory that closes over static n and L_vec (so shapes are constant).
 532    """
 533    L_vec = jnp.asarray(L_vec)
 534    @jax.jit  # optional; multiple_cheb_fast is already jitted
 535    def mcheb(x):
 536        x = jnp.asarray(x)
 537        return multiple_cheb_fast(x, L_vec, n)
 538    return mcheb
 539
 540
 541#Initialize fully connected neyral network Return the initial parameters and the function for the forward pass
 542def fconNN(width,activation = jax.nn.tanh,key = 0,mlp = False,ftype = None,fargs = None,static = None):
 543    """
 544    Initialize fully connected neural network
 545    ----------
 546    Parameters
 547    ----------
 548    width : list
 549
 550        List with the layers width
 551
 552    activation : jax.nn activation
 553
 554        The activation function. Default jax.nn.tanh
 555
 556    key : int
 557
 558        Seed for parameters initialization. Default 0
 559
 560    mlp : logical
 561
 562        Whether to consider a modified multilayer perceptron. Assumes all hidden layers have the same dimension.
 563
 564    ftype : str
 565
 566        Type of feature transformation to use: None, 'ff', 'daff','daff_bias', 'cheb', 'cheb_bias'.
 567
 568    fargs : list
 569
 570        Arguments for deature transformation:
 571
 572        For 'ff': A list with the number of frequences and value of greatest frequence standard deviation.
 573
 574        For 'daff' and 'daff' bias: A dicitionary with a list with the size of rectangles and the type of boundary condition. If its a list, than boundary conditions is dirichlet.
 575
 576    static : function
 577
 578        A static function to sum to the neural network output.
 579
 580    Returns
 581    -------
 582    dict with initial parameters and the function for the forward pass
 583    """
 584    #Initialize parameters with Glorot initialization
 585    initializer = jax.nn.initializers.glorot_normal()
 586    params = list()
 587    if static is None:
 588        static = lambda x: 0.0
 589
 590    #Feature mapping
 591    if ftype == 'ff': #Fourrier features
 592        for s in range(fargs[0]):
 593            sd = fargs[1] ** ((s + 1)/fargs[0])
 594            if s == 0:
 595                Bff = sd*jax.random.normal(jax.random.PRNGKey(key + s + 1),(width[0],int(width[1]/2)))
 596            else:
 597                Bff = jnp.append(Bff,sd*jax.random.normal(jax.random.PRNGKey(key + s + 1),(width[0],int(width[1]/2))),1)
 598        @jax.jit
 599        def phi(x):
 600            x = x @ Bff
 601            return jnp.concatenate([jnp.sin(2 * jnp.pi * x),jnp.cos(2 * jnp.pi * x)],axis = -1)
 602        width = width[1:]
 603        width[0] = 2*Bff.shape[1]
 604    elif ftype == 'daff' or ftype == 'daff_bias':
 605        if not isinstance(fargs, dict):
 606            fargs = {'L': fargs,'bc': "dirichlet"}
 607        phi,lamb = multiple_daff(list(fargs.values())[0],kmax_per_axis = [width[1]] * width[0],bc = list(fargs.values())[1])
 608        width = width[1:]
 609        width[0] = lamb.shape[0]
 610    elif ftype == 'cheb' or ftype == 'cheb_bias':
 611        phi = multiple_cheb(fargs,n = width[1])
 612        width = width[1:]
 613        width[0] = len(fargs)*width[0]
 614    else:
 615        @jax.jit
 616        def phi(x):
 617            return x
 618
 619    #Initialize parameters
 620    if mlp:
 621        k = jax.random.split(jax.random.PRNGKey(key),4)
 622        WU = initializer(k[0],(width[0],width[1]),jnp.float32)
 623        BU = initializer(k[1],(1,width[1]),jnp.float32)
 624        WV = initializer(k[2],(width[0],width[1]),jnp.float32)
 625        BV = initializer(k[3],(1,width[1]),jnp.float32)
 626        params.append({'WU':WU,'BU':BU,'WV':WV,'BV':BV})
 627    key = jax.random.split(jax.random.PRNGKey(key + 1),len(width)-1) #Seed for initialization
 628    for key,lin,lout in zip(key,width[:-1],width[1:]):
 629        W = initializer(key,(lin,lout),jnp.float32)
 630        B = initializer(key,(1,lout),jnp.float32)
 631        params.append({'W':W,'B':B})
 632
 633    #Define function for forward pass
 634    if mlp:
 635        if ftype != 'daff' and ftype != 'cheb':
 636            @jax.jit
 637            def forward(x,params):
 638                encode,*hidden,output = params
 639                sx = static(x)
 640                x = phi(x)
 641                U = activation(x @ encode['WU'] + encode['BU'])
 642                V = activation(x @ encode['WV'] + encode['BV'])
 643                for layer in hidden:
 644                    x = activation(x @ layer['W'] + layer['B'])
 645                    x = x * U + (1 - x) * V
 646                return x @ output['W'] + output['B'] + sx
 647        else:
 648            @jax.jit
 649            def forward(x,params):
 650                encode,*hidden,output = params
 651                sx = static(x)
 652                x = phi(x)
 653                U = activation(x @ encode['WU'])
 654                V = activation(x @ encode['WV'])
 655                for layer in hidden:
 656                    x = activation(x @ layer['W'])
 657                    x = x * U + (1 - x) * V
 658                return x @ output['W'] + sx
 659    else:
 660        if ftype != 'daff' and ftype != 'cheb':
 661            @jax.jit
 662            def forward(x,params):
 663                *hidden,output = params
 664                sx = static(x)
 665                x = phi(x)
 666                for layer in hidden:
 667                    x = activation(x @ layer['W'] + layer['B'])
 668                return x @ output['W'] + output['B'] + sx
 669        else:
 670            @jax.jit
 671            def forward(x,params):
 672                *hidden,output = params
 673                sx = static(x)
 674                x = phi(x)
 675                for layer in hidden:
 676                    x = activation(x @ layer['W'])
 677                return x @ output['W'] + sx
 678
 679    #Return initial parameters and forward function
 680    return {'params': params,'forward': forward}
 681
 682#Get activation from string
 683def get_activation(act):
 684    """
 685    Return activation function from string
 686    ----------
 687    Parameters
 688    ----------
 689    act : str
 690
 691        Name of the activation function. Default 'tanh'
 692
 693    Returns
 694    -------
 695    jax.nn activation function
 696    """
 697    if act == 'tanh':
 698        return jax.nn.tanh
 699    elif act == 'relu':
 700        return jax.nn.relu
 701    elif act == 'relu6':
 702        return jax.nn.relu6
 703    elif act == 'sigmoid':
 704        return jax.nn.sigmoid
 705    elif act == 'softplus':
 706        return jax.nn.softplus
 707    elif act == 'sparse_plus':
 708        return jx.nn.sparse_plus
 709    elif act == 'soft_sign':
 710        return jax.nn.soft_sign
 711    elif act == 'silu':
 712        return jax.nn.silu
 713    elif act == 'swish':
 714        return jax.nn.swish
 715    elif act == 'log_sigmoid':
 716        return jax.nn.log_sigmoid
 717    elif act == 'leaky_relu':
 718        return jax.nn.leaky_relu
 719    elif act == 'hard_sigmoid':
 720        return jax.nn.hard_sigmoid
 721    elif act == 'hard_silu':
 722        return jax.nn.hard_silu
 723    elif act == 'hard_swish':
 724        return jax.nn.hard_swish
 725    elif act == 'hard_tanh':
 726        return jax.nn.hard_tanh
 727    elif act == 'elu':
 728        return jax.nn.elu
 729    elif act == 'celu':
 730        return jax.nn.celu
 731    elif act == 'selu':
 732        return jax.nn.selu
 733    elif act == 'gelu':
 734        return jax.nn.gelu
 735    elif act == 'glu':
 736        return jax.nn.glu
 737    elif act == 'squareplus':
 738        return  jax.nn.squareplus
 739    elif act == 'mish':
 740        return jax.nn.mish
 741
 742#Training PINN
 743def train_PINN(data,width,pde,test_data = None,epochs = 100,at_each = 10,activation = 'tanh',neumann = False,oper_neumann = False,sa = False,c = {'ws': 1,'wr': 1,'w0': 100,'wb': 1},inverse = False,initial_par = None,lr = 0.001,b1 = 0.9,b2 = 0.999,eps = 1e-08,eps_root = 0.0,key = 0,epoch_print = 100,save = False,file_name = 'result_pinn',exp_decay = False,transition_steps = 1000,decay_rate = 0.9,mlp = False):
 744    """
 745    Train a Physics-informed Neural Network
 746    ----------
 747    Parameters
 748    ----------
 749    data : dict
 750
 751        Data generated by the jinnax.data.generate_PINNdata function
 752
 753    width : list
 754
 755        A list with the width of each layer
 756
 757    pde : function
 758
 759        The partial differential operator. Its arguments are u, x and t
 760
 761    test_data : dict, None
 762
 763        A dictionay with test data for L2 error calculation generated by the jinnax.data.generate_PINNdata function. Default None for not calculating L2 error
 764
 765    epochs : int
 766
 767        Number of training epochs. Default 100
 768
 769    at_each : int
 770
 771        Save results for epochs multiple of at_each. Default 10
 772
 773    activation : str
 774
 775        The name of the activation function of the neural network. Default 'tanh'
 776
 777    neumann : logical
 778
 779        Whether to consider Neumann boundary conditions
 780
 781    oper_neumann : function
 782
 783        Penalization of Neumann boundary conditions
 784
 785    sa : logical
 786
 787        Whether to consider self-adaptative PINN
 788
 789    c : dict
 790
 791        Dictionary with the hyperparameters of the self-adaptative sigmoid mask for the initial (w0), sensor (ws) and collocation (wr) points. The weights of the boundary points is fixed to 1
 792
 793    inverse : logical
 794
 795        Whether to estimate parameters of the PDE
 796
 797    initial_par : jax.numpy.array
 798
 799        Initial value of the parameters of the PDE in an inverse problem
 800
 801    lr,b1,b2,eps,eps_root: float
 802
 803        Hyperparameters of the Adam algorithm. Default lr = 0.001, b1 = 0.9, b2 = 0.999, eps = 1e-08, eps_root = 0.0
 804
 805    key : int
 806
 807        Seed for parameters initialization. Default 0
 808
 809    epoch_print : int
 810
 811        Number of epochs to calculate and print test errors. Default 100
 812
 813    save : logical
 814
 815        Whether to save the current parameters. Default False
 816
 817    file_name : str
 818
 819        File prefix to save the current parameters. Default 'result_pinn'
 820
 821    exp_decay : logical
 822
 823        Whether to consider exponential decay of learning rate. Default False
 824
 825    transition_steps : int
 826
 827        Number of steps for exponential decay. Default 1000
 828
 829    decay_rate : float
 830
 831        Rate of exponential decay. Default 0.9
 832
 833    mlp : logical
 834
 835        Whether to consider modifed multi-layer perceptron
 836
 837    Returns
 838    -------
 839    dict-like object with the estimated function, the estimated parameters, the neural network function for the forward pass and the training time
 840    """
 841
 842    #Initialize architecture
 843    nnet = fconNN(width,get_activation(activation),key,mlp)
 844    forward = nnet['forward']
 845
 846    #Initialize self adaptative weights
 847    par_sa = {}
 848    if sa:
 849        #Initialize wheights close to zero
 850        ksa = jax.random.randint(jax.random.PRNGKey(key),(5,),1,1000000)
 851        if data['sensor'] is not None:
 852            par_sa.update({'ws': c['ws'] * jax.random.uniform(key = jax.random.PRNGKey(ksa[0]),shape = (data['sensor'].shape[0],1))})
 853        if data['initial'] is not None:
 854            par_sa.update({'w0': c['w0'] * jax.random.uniform(key = jax.random.PRNGKey(ksa[1]),shape = (data['initial'].shape[0],1))})
 855        if data['collocation'] is not None:
 856            par_sa.update({'wr': c['wr'] * jax.random.uniform(key = jax.random.PRNGKey(ksa[2]),shape = (data['collocation'].shape[0],1))})
 857        if data['boundary'] is not None:
 858            par_sa.update({'wb': c['wr'] * jax.random.uniform(key = jax.random.PRNGKey(ksa[3]),shape = (data['boundary'].shape[0],1))})
 859
 860    #Store all parameters
 861    params = {'net': nnet['params'],'inverse': initial_par,'sa': par_sa}
 862
 863    #Save config file
 864    if save:
 865        pickle.dump({'train_data': data,'epochs': epochs,'activation': activation,'init_params': params,'forward': forward,'width': width,'pde': pde,'lr': lr,'b1': b1,'b2': b2,'eps': eps,'eps_root': eps_root,'key': key,'inverse': inverse,'sa': sa},open(file_name + '_config.pickle','wb'), protocol = pickle.HIGHEST_PROTOCOL)
 866
 867    #Define loss function
 868    if sa:
 869        #Define loss function
 870        @jax.jit
 871        def lf(params,x):
 872            loss = 0
 873            if x['sensor'] is not None:
 874                #Term that refers to sensor data
 875                loss = loss + jnp.mean(MSE_SA(forward(x['sensor'],params['net']),x['usensor'],params['sa']['ws']))
 876            if x['boundary'] is not None:
 877                if neumann:
 878                    #Neumann coditions
 879                    xb = x['boundary'][:,:-1].reshape((x['boundary'].shape[0],x['boundary'].shape[1] - 1))
 880                    tb = x['boundary'][:,-1].reshape((x['boundary'].shape[0],1))
 881                    loss = loss + jnp.mean(oper_neumann(lambda x,t: forward(jnp.append(x,t,1),params['net']),xb,tb,params['sa']['wb']))
 882                else:
 883                    #Term that refers to boundary data
 884                    loss = loss + jnp.mean(MSE_SA(forward(x['boundary'],params['net']),x['uboundary'],params['sa']['wb']))
 885            if x['initial'] is not None:
 886                #Term that refers to initial data
 887                loss = loss + jnp.mean(MSE_SA(forward(x['initial'],params['net']),x['uinitial'],params['sa']['w0']))
 888            if x['collocation'] is not None:
 889                #Term that refers to collocation points
 890                x_col = x['collocation'][:,:-1].reshape((x['collocation'].shape[0],x['collocation'].shape[1] - 1))
 891                t_col = x['collocation'][:,-1].reshape((x['collocation'].shape[0],1))
 892                if inverse:
 893                    loss = loss + jnp.mean(MSE_SA(pde(lambda x,t: forward(jnp.append(x,t,1),params['net']),x_col,t_col,params['inverse']),0,params['sa']['wr']))
 894                else:
 895                    loss = loss + jnp.mean(MSE_SA(pde(lambda x,t: forward(jnp.append(x,t,1),params['net']),x_col,t_col),0,params['sa']['wr']))
 896            return loss
 897    else:
 898        @jax.jit
 899        def lf(params,x):
 900            loss = 0
 901            if x['sensor'] is not None:
 902                #Term that refers to sensor data
 903                loss = loss + jnp.mean(MSE(forward(x['sensor'],params['net']),x['usensor']))
 904            if x['boundary'] is not None:
 905                if neumann:
 906                    #Neumann coditions
 907                    xb = x['boundary'][:,:-1].reshape((x['boundary'].shape[0],x['boundary'].shape[1] - 1))
 908                    tb = x['boundary'][:,-1].reshape((x['boundary'].shape[0],1))
 909                    loss = loss + jnp.mean(oper_neumann(lambda x,t: forward(jnp.append(x,t,1),params['net']),xb,tb))
 910                else:
 911                    #Term that refers to boundary data
 912                    loss = loss + jnp.mean(MSE(forward(x['boundary'],params['net']),x['uboundary']))
 913            if x['initial'] is not None:
 914                #Term that refers to initial data
 915                loss = loss + jnp.mean(MSE(forward(x['initial'],params['net']),x['uinitial']))
 916            if x['collocation'] is not None:
 917                #Term that refers to collocation points
 918                x_col = x['collocation'][:,:-1].reshape((x['collocation'].shape[0],x['collocation'].shape[1] - 1))
 919                t_col = x['collocation'][:,-1].reshape((x['collocation'].shape[0],1))
 920                if inverse:
 921                    loss = loss + jnp.mean(MSE(pde(lambda x,t: forward(jnp.append(x,t,1),params['net']),x_col,t_col,params['inverse']),0))
 922                else:
 923                    loss = loss + jnp.mean(MSE(pde(lambda x,t: forward(jnp.append(x,t,1),params['net']),x_col,t_col),0))
 924            return loss
 925
 926    #Initialize Adam Optmizer
 927    if exp_decay:
 928        lr = optax.exponential_decay(lr,transition_steps,decay_rate)
 929    optimizer = optax.adam(lr,b1,b2,eps,eps_root)
 930    opt_state = optimizer.init(params)
 931
 932    #Define the gradient function
 933    grad_loss = jax.jit(jax.grad(lf,0))
 934
 935    #Define update function
 936    @jax.jit
 937    def update(opt_state,params,x):
 938        #Compute gradient
 939        grads = grad_loss(params,x)
 940        #Invert gradient of self-adaptative wheights
 941        if sa:
 942            for w in grads['sa']:
 943                grads['sa'][w] = - grads['sa'][w]
 944        #Calculate parameters updates
 945        updates, opt_state = optimizer.update(grads, opt_state)
 946        #Update parameters
 947        params = optax.apply_updates(params, updates)
 948        #Return state of optmizer and updated parameters
 949        return opt_state,params
 950
 951    ###Training###
 952    t0 = time.time()
 953    #Initialize alive_bar for tracing in terminal
 954    with alive_bar(epochs) as bar:
 955        #For each epoch
 956        for e in range(epochs):
 957            #Update optimizer state and parameters
 958            opt_state,params = update(opt_state,params,data)
 959            #After epoch_print epochs
 960            if e % epoch_print == 0:
 961                #Compute elapsed time and current error
 962                l = 'Time: ' + str(round(time.time() - t0)) + ' s Loss: ' + str(jnp.round(lf(params,data),6))
 963                #If there is test data, compute current L2 error
 964                if test_data is not None:
 965                    #Compute L2 error
 966                    l2_test = L2error(forward(test_data['xt'],params['net']),test_data['u']).tolist()
 967                    l = l + ' L2 error: ' + str(jnp.round(l2_test,3))
 968                if inverse:
 969                    l = l + ' Parameter: ' + str(jnp.round(params['inverse'].tolist(),6))
 970                #Print
 971                print(l)
 972            if ((e % at_each == 0 and at_each != epochs) or e == epochs - 1) and save:
 973                #Save current parameters
 974                pickle.dump({'params': params,'width': width,'time': time.time() - t0,'loss': lf(params,data)},open(file_name + '_epoch' + str(e).rjust(6, '0') + '.pickle','wb'), protocol = pickle.HIGHEST_PROTOCOL)
 975            #Update alive_bar
 976            bar()
 977    #Define estimated function
 978    def u(xt):
 979        return forward(xt,params['net'])
 980
 981    return {'u': u,'params': params,'forward': forward,'time': time.time() - t0}
 982
 983#Training PINN
 984def train_Matern_PINN(data,width,pde,test_data = None,params = None,d = 2,N = 128,L = 1,alpha = 1,kappa = 1,sigma = 100,bsize = 1024,resample = False,epochs = 100,at_each = 10,activation = 'tanh',
 985    neumann = False,oper_neumann = None,inverse = False,initial_par = None,lr = 0.001,b1 = 0.9,b2 = 0.999,eps = 1e-08,eps_root = 0.0,key = 0,epoch_print = 1,save = False,file_name = 'result_pinn',
 986    exp_decay = True,transition_steps = 100,decay_rate = 0.9,mlp = True,ftype = None,fargs = None,q = 4,w = None,periodic = False,static = None,opt = 'LBFGS'):
 987    """
 988    Train a Physics-informed Neural Network
 989    ----------
 990    Parameters
 991    ----------
 992    data : dict
 993
 994        Data generated by the jinnax.data.generate_PINNdata function
 995
 996    width : list
 997
 998        A list with the width of each layer
 999
1000    pde : function
1001
1002        The partial differential operator. Its arguments are u, x and t
1003
1004    test_data : dict, None
1005
1006        A dictionay with test data for L2 error calculation generated by the jinnax.data.generate_PINNdata function. Default None for not calculating L2 error
1007
1008    params : list
1009
1010        Initial parameters for the neural network. Default None to initialize randomly
1011
1012    d : int
1013
1014        Dimension of the problem including the time variable if present. Default 2
1015
1016    N : int
1017
1018        Size of grid in each dimension. Default 128
1019
1020    L : list of float
1021
1022        The domain of the function in each coordinate is [0,L[1]]. If a float, repeat the same interval for all coordinates. Default 1
1023
1024    kappa,alpha,sigma : float
1025
1026        Parameters of the Matern process
1027
1028    bsize : int
1029
1030        Batch size for weak norm computation. Default 1024
1031
1032    resample : logical
1033
1034        Whether to resample the test functions at each epoch
1035
1036    epochs : int
1037
1038        Number of training epochs. Default 100
1039
1040    at_each : int
1041
1042        Save results for epochs multiple of at_each. Default 10
1043
1044    activation : str
1045
1046        The name of the activation function of the neural network. Default 'tanh'
1047
1048    neumann : logical
1049
1050        Whether to consider Neumann boundary conditions
1051
1052    oper_neumann : function
1053
1054        Penalization of Neumann boundary conditions
1055
1056    inverse : logical
1057
1058        Whether to estimate parameters of the PDE
1059
1060    initial_par : jax.numpy.array
1061
1062        Initial value of the parameters of the PDE in an inverse problem
1063
1064    lr,b1,b2,eps,eps_root: float
1065
1066        Hyperparameters of the Adam algorithm. Default lr = 0.001, b1 = 0.9, b2 = 0.999, eps = 1e-08, eps_root = 0.0
1067
1068    key : int
1069
1070        Seed for parameters initialization. Default 0
1071
1072    epoch_print : int
1073
1074        Number of epochs to calculate and print test errors. Default 1
1075
1076    save : logical
1077
1078        Whether to save the current parameters. Default False
1079
1080    file_name : str
1081
1082        File prefix to save the current parameters. Default 'result_pinn'
1083
1084    exp_decay : logical
1085
1086        Whether to consider exponential decay of learning rate. Default True
1087
1088    transition_steps : int
1089
1090        Number of steps for exponential decay. Default 100
1091
1092    decay_rate : float
1093
1094        Rate of exponential decay. Default 0.9
1095
1096    mlp : logical
1097
1098        Whether to consider modifed multilayer perceptron
1099
1100    ftype : str
1101
1102        Type of feature transformation to use: None, 'ff', 'daff','daff_bias', 'cheb', 'cheb_bias'.
1103
1104    fargs : list
1105
1106        Arguments for deature transformation:
1107
1108        For 'ff': A list with the number of frequences and value of greatest frequence standard deviation.
1109
1110        For 'daff' and 'daff' bias: A dicitionary with a list with the size of rectangles and the type of boundary condition. If its a list, than boundary conditions is dirichlet.
1111
1112    q : int
1113
1114        Power of weights mask. Default 4
1115
1116    w : dict
1117
1118        Initila weights for self-adaptive scheme.
1119
1120    periodic : logical
1121
1122        Whether to consider periodic test functions. Default False.
1123
1124    static : function
1125
1126        A static function to sum to the neural network output.
1127
1128    opt : str
1129
1130        Optimizer. Default LBFGS.
1131
1132    Returns
1133    -------
1134    dict-like object with the estimated function, the estimated parameters, the neural network function for the forward pass and the loss, L2error and training time at each epoch
1135    """
1136    #Initialize architecture
1137    nnet = fconNN(width,get_activation(activation),key,mlp,ftype,fargs,static)
1138    forward = nnet['forward']
1139    if params is not None:
1140        nnet['params'] = params
1141
1142    #Generate from Matern process
1143    if sigma > 0:
1144        if isinstance(L,float) or isinstance(L,int):
1145            L = d*[L]
1146        #Grid for weak norm
1147        grid = [jnp.linspace(0,L[i],N) for i in range(d)]
1148        grid = jnp.meshgrid(*grid, indexing='ij')
1149        grid = jnp.stack(grid, axis=-1).reshape((-1, d))
1150        #Set sigma
1151        if data['boundary'] is not None:
1152            gen = generate_matern_sample_batch(d = d,N = N,L = L,kappa = kappa,alpha = alpha,sigma = sigma)
1153            tf = gen(jax.random.split(jax.random.PRNGKey(key + 1),(bsize,))[:,0])
1154            if neumann:
1155                loss_boundary = oper_neumann(lambda x: forward(x,params['net']),data['boundary'])
1156            else:
1157                loss_boundary = jnp.mean(MSE(forward(data['boundary'],nnet['params']),data['uboundary']))
1158            output_w = pde(lambda x: forward(x,nnet['params']),grid)
1159            integralOmega = jax.vmap(lambda psi: jnp.mean(psi*output_w.reshape((N,) * d)))(tf)
1160            loss_res_weak = jnp.mean(integralOmega ** 2)
1161            sigma = float(jnp.sqrt(loss_boundary/loss_res_weak).tolist())
1162            del gen
1163            gen = generate_matern_sample_batch(d = d,N = N,L = L,kappa = kappa,alpha = alpha,sigma = sigma,periodic = periodic)
1164            tf = sigma*tf
1165        else:
1166            gen = generate_matern_sample_batch(d = d,N = N,L = L,kappa = kappa,alpha = alpha,sigma = sigma,periodic = periodic)
1167            tf = gen(jax.random.split(jax.random.PRNGKey(key + 1),(bsize,))[:,0])
1168
1169    #Define loss function
1170    @jax.jit
1171    def lf_each(params,x,k):
1172        if sigma > 0:
1173            #Term that refers to weak loss
1174            if resample:
1175                test_functions = gen(jax.random.split(jax.random.PRNGKey(k[0]),(bsize,))[:,0])
1176            else:
1177                test_functions = tf
1178        loss_sensor = loss_boundary = loss_initial = loss_res = loss_res_weak = 0
1179        if x['sensor'] is not None:
1180            #Term that refers to sensor data
1181            loss_sensor = jnp.mean(MSE(forward(x['sensor'],params['net']),x['usensor']))
1182        if x['boundary'] is not None:
1183            if neumann:
1184                #Neumann coditions
1185                loss_boundary = oper_neumann(lambda x: forward(x,params['net']),x['boundary'])
1186            else:
1187                #Term that refers to boundary data
1188                loss_boundary = MSE(forward(x['boundary'],params['net']),x['uboundary'])
1189        if x['initial'] is not None:
1190            #Term that refers to initial data
1191            loss_initial = MSE(forward(x['initial'],params['net']),x['uinitial'])
1192        if x['collocation'] is not None and sigma == 0:
1193            if inverse:
1194                output = pde(lambda x: forward(x,params['net']),x['collocation'],params['inverse'])
1195                loss_res = MSE(output,0)
1196            else:
1197                output = pde(lambda x: forward(x,params['net']),x['collocation'])
1198                loss_res = MSE(output,0)
1199        if sigma > 0:
1200            #Term that refers to weak loss
1201            if inverse:
1202                output_w = pde(lambda x: forward(x,params['net']),grid,params['inverse'])
1203                integralOmega = jax.vmap(lambda psi: jnp.mean(psi*output_w.reshape((N,) * d)))(test_functions)
1204                loss_res_weak = jnp.mean(integralOmega ** 2)
1205            else:
1206                output_w = pde(lambda x: forward(x,params['net']),grid)
1207                integralOmega = jax.vmap(lambda psi: jnp.mean(psi*output_w.reshape((N,) * d)))(test_functions)
1208                loss_res_weak = jnp.mean(integralOmega ** 2)
1209        return {'ls': loss_sensor,'lb': loss_boundary,'li': loss_initial,'lc': loss_res,'lc_weak': loss_res_weak}
1210
1211    @jax.jit
1212    def lf(params,x,k):
1213        l = lf_each(params,x,k)
1214        w = params['w']
1215        loss = jnp.mean((w['ws'] ** q)*l['ls']) + jnp.mean((w['wb'] ** q)*l['lb']) + jnp.mean((w['wi'] ** q)*l['li']) + jnp.mean((w['wc'] ** q)*l['lc']) + (w['wc_weak'] ** q)*l['lc_weak']
1216        if opt != 'LBFGS':
1217            return loss
1218        else:
1219            l2 = None
1220            if test_data is not None:
1221                l2 = L2error(forward(test_data['sensor'],params['net']),test_data['usensor'])
1222            return loss,{'loss': loss,'l2': l2}
1223
1224    #Initialize self-adaptive weights
1225    if w is None:
1226        w = {'ws': jnp.array(1.0),'wb': jnp.array(1.0),'wi': jnp.array(1.0),'wc': jnp.array(1.0),'wc_weak': jnp.array(1.0)}
1227    if q != 0:
1228        if data['sensor'] is not None:
1229            w['ws'] = w['ws'] + 0.05*jax.random.normal(jax.random.PRNGKey(key+1),(data['sensor'].shape[0],1))
1230        if data['boundary'] is not None:
1231            w['wb'] = w['wb'] + 0.05*jax.random.normal(jax.random.PRNGKey(key+2),(data['boundary'].shape[0],1))
1232        if data['initial'] is not None:
1233            w['wi'] = w['wi'] + 0.05*jax.random.normal(jax.random.PRNGKey(key+3),(data['initial'].shape[0],1))
1234        if data['collocation'] is not None:
1235            w['wc'] = w['wc'] + 0.05*jax.random.normal(jax.random.PRNGKey(key+4),(data['collocation'].shape[0],1))
1236
1237    #Store all parameters
1238    params = {'net': nnet['params'],'inverse': initial_par,'w': w}
1239
1240    #Save config file
1241    if save:
1242        pickle.dump({'train_data': data,'epochs': epochs,'activation': activation,'init_params': params,'forward': forward,'width': width,'pde': pde,'lr': lr,'b1': b1,'b2': b2,'eps': eps,'eps_root': eps_root,'key': key,'inverse': inverse},open(file_name + '_config.pickle','wb'), protocol = pickle.HIGHEST_PROTOCOL)
1243
1244    #Initialize Adam Optmizer
1245    if opt != 'LBFGS':
1246        print('--------- GRADIENT DESCENT OPTIMIZER ---------')
1247        if exp_decay:
1248            lr = optax.exponential_decay(lr,transition_steps,decay_rate)
1249        optimizer = optax.adam(lr,b1,b2,eps,eps_root)
1250        opt_state = optimizer.init(params)
1251
1252        #Define the gradient function
1253        grad_loss = jax.jit(jax.grad(lf,0))
1254
1255        #Define update function
1256        @jax.jit
1257        def update(opt_state,params,x,k):
1258            #Compute gradient
1259            grads = grad_loss(params,x,k)
1260            #Calculate parameters updates
1261            updates, opt_state = optimizer.update(grads, opt_state)
1262            #Update parameters
1263            if q != 0:
1264                updates = {**updates, 'w': jax.tree_util.tree_map(lambda x: -x, updates['w'])} #Change signs of weights
1265            params = optax.apply_updates(params, updates)
1266            #Return state of optmizer and updated parameters
1267            return opt_state,params
1268    else:
1269        print('--------- LBFGS OPTIMIZER ---------')
1270        @jax.jit
1271        def loss_LBFGS(params):
1272            return lf(params,data,key + 234)
1273        solver = LBFGS(fun = loss_LBFGS,has_aux = True,maxiter = epochs,tol = 1e-9,verbose = False,linesearch = 'zoom',history_size = 100)  # linesearch='zoom' by default
1274        state = solver.init_state(params)
1275
1276    ###Training###
1277    t0 = time.time()
1278    k = jax.random.split(jax.random.PRNGKey(key+234),(epochs,))
1279    sloss = []
1280    sL2 = []
1281    stime = []
1282    #Initialize alive_bar for tracing in terminal
1283    with alive_bar(epochs) as bar:
1284        #For each epoch
1285        for e in range(epochs):
1286            if opt != 'LBFGS':
1287                #Update optimizer state and parameters
1288                opt_state,params = update(opt_state,params,data,k[e,:])
1289                sloss.append(lf(params,data,k[e,:]))
1290                if test_data is not None:
1291                    sL2.append(L2error(forward(test_data['sensor'],params['net']),test_data['usensor']))
1292            else:
1293                params, state = solver.update(params, state)
1294                sL2.append(state.aux["l2"])
1295                sloss.append(state.aux["loss"])
1296            stime.append(time.time() - t0)
1297            #After epoch_print epochs
1298            if e % epoch_print == 0:
1299                #Compute elapsed time and current error
1300                l = 'Time: ' + str(round(time.time() - t0)) + ' s Loss: ' + str(jnp.round(sloss[-1],6))
1301                #If there is test data, compute current L2 error
1302                if test_data is not None:
1303                    #Compute L2 error
1304                    l = l + ' L2 error: ' + str(jnp.round(sL2[-1],6))
1305                if inverse:
1306                    l = l + ' Parameter: ' + str(jnp.round(params['inverse'].tolist(),6))
1307                #Print
1308                print(l)
1309            if ((e % at_each == 0 and at_each != epochs) or e == epochs - 1) and save:
1310                #Save current parameters
1311                pickle.dump({'params': params,'width': width,'time': stime,'loss': sloss,'L2error': sL2},open(file_name + '_epoch' + str(e).rjust(6, '0') + '.pickle','wb'), protocol = pickle.HIGHEST_PROTOCOL)
1312            #Update alive_bar
1313            bar()
1314    #Define estimated function
1315    def u(xt):
1316        return forward(xt,params['net'])
1317
1318    return {'u': u,'params': params,'forward': forward,'time': time.time() - t0,'loss_each': lf_each(params,data,[key + 100]),'loss': sloss,'L2error': sL2}
1319
1320
1321#Process result
1322def process_result(test_data,fit,train_data,plot = True,plot_test = True,times = 5,d2 = True,save = False,show = True,file_name = 'result_pinn',print_res = True,p = 1):
1323    """
1324    Process the results of a Physics-informed Neural Network
1325    ----------
1326
1327    Parameters
1328    ----------
1329    test_data : dict
1330
1331        A dictionay with test data for L2 error calculation generated by the jinnax.data.generate_PINNdata function
1332
1333    fit : function
1334
1335        The fitted function
1336
1337    train_data : dict
1338
1339        Training data generated by the jinnax.data.generate_PINNdata
1340
1341    plot : logical
1342
1343        Whether to generate plots comparing the exact and estimated solutions when the spatial dimension is one. Default True
1344
1345    plot_test : logical
1346
1347        Whether to plot the test data. Default True
1348
1349    times : int
1350
1351        Number of points along the time interval to plot. Default 5
1352
1353    d2 : logical
1354
1355        Whether to plot 2D plot when the spatial dimension is one. Default True
1356
1357    save : logical
1358
1359        Whether to save the plots. Default False
1360
1361    show : logical
1362
1363        Whether to show the plots. Default True
1364
1365    file_name : str
1366
1367        File prefix to save the plots. Default 'result_pinn'
1368
1369    print_res : logical
1370
1371        Whether to print the L2 error. Default True
1372
1373    p : int
1374
1375        Output dimension. Default 1
1376
1377    Returns
1378    -------
1379    pandas data frame with L2 and MSE errors
1380    """
1381
1382    #Dimension
1383    d = test_data['xt'].shape[1] - 1
1384
1385    #Number of plots multiple of 5
1386    times = 5 * round(times/5.0)
1387
1388    #Data
1389    td = get_train_data(train_data)
1390    xt_train = td['x']
1391    u_train = td['y']
1392    upred_train = fit(xt_train)
1393    upred_test = fit(test_data['xt'])
1394
1395    #Results
1396    l2_error_test = L2error(upred_test,test_data['u']).tolist()
1397    MSE_test = jnp.mean(MSE(upred_test,test_data['u'])).tolist()
1398    l2_error_train = L2error(upred_train,u_train).tolist()
1399    MSE_train = jnp.mean(MSE(upred_train,u_train)).tolist()
1400
1401    df = pd.DataFrame(np.array([l2_error_test,MSE_test,l2_error_train,MSE_train]).reshape((1,4)),
1402        columns=['l2_error_test','MSE_test','l2_error_train','MSE_train'])
1403    if print_res:
1404        print('L2 error test: ' + str(jnp.round(l2_error_test,6)) + ' L2 error train: ' + str(jnp.round(l2_error_train,6)) + ' MSE error test: ' + str(jnp.round(MSE_test,6)) + ' MSE error train: ' + str(jnp.round(MSE_train,6)) )
1405
1406    #Plots
1407    if d == 1 and p ==1 and plot:
1408        plot_pinn1D(times,test_data['xt'],test_data['u'],upred_test,d2,save,show,file_name)
1409    elif p == 2 and plot:
1410        plot_pinn_out2D(times,test_data['xt'],test_data['u'],upred_test,save,show,file_name,plot_test)
1411
1412    return df
1413
1414#Plot results for d = 1
1415def plot_pinn1D(times,xt,u,upred,d2 = True,save = False,show = True,file_name = 'result_pinn',title_1d = '',title_2d = ''):
1416    """
1417    Plot the prediction of a 1D PINN
1418    ----------
1419
1420    Parameters
1421    ----------
1422    times : int
1423
1424        Number of points along the time interval to plot. Default 5
1425
1426    xt : jax.numpy.array
1427
1428        Test data xt array
1429
1430    u : jax.numpy.array
1431
1432        Test data u(x,t) array
1433
1434    upred : jax.numpy.array
1435
1436        Predicted upred(x,t) array on test data
1437
1438    d2 : logical
1439
1440        Whether to plot 2D plot. Default True
1441
1442    save : logical
1443
1444        Whether to save the plots. Default False
1445
1446    show : logical
1447
1448        Whether to show the plots. Default True
1449
1450    file_name : str
1451
1452        File prefix to save the plots. Default 'result_pinn'
1453
1454    title_1d : str
1455
1456        Title of 1D plot
1457
1458    title_2d : str
1459
1460        Title of 2D plot
1461
1462    Returns
1463    -------
1464    None
1465    """
1466    #Initialize
1467    fig, ax = plt.subplots(int(times/5),5,figsize = (10*int(times/5),3*int(times/5)))
1468    tlo = jnp.min(xt[:,-1])
1469    tup = jnp.max(xt[:,-1])
1470    ylo = jnp.min(u)
1471    ylo = ylo - 0.1*jnp.abs(ylo)
1472    yup = jnp.max(u)
1473    yup = yup + 0.1*jnp.abs(yup)
1474    k = 0
1475    t_values = np.linspace(tlo,tup,times)
1476
1477    #Create
1478    for i in range(int(times/5)):
1479        for j in range(5):
1480            if k < len(t_values):
1481                t = t_values[k]
1482                t = xt[jnp.abs(xt[:,-1] - t) == jnp.min(jnp.abs(xt[:,-1] - t)),-1][0].tolist()
1483                x_plot = xt[xt[:,-1] == t,:-1]
1484                y_plot = upred[xt[:,-1] == t,:]
1485                u_plot = u[xt[:,-1] == t,:]
1486                if int(times/5) > 1:
1487                    ax[i,j].plot(x_plot[:,0],u_plot[:,0],'b-',linewidth=2,label='Exact')
1488                    ax[i,j].plot(x_plot[:,0],y_plot,'r--',linewidth=2,label='Prediction')
1489                    ax[i,j].set_title('$t = %.2f$' % (t),fontsize=10)
1490                    ax[i,j].set_xlabel(' ')
1491                    ax[i,j].set_ylim([1.3 * ylo.tolist(),1.3 * yup.tolist()])
1492                else:
1493                    ax[j].plot(x_plot[:,0],u_plot[:,0],'b-',linewidth=2,label='Exact')
1494                    ax[j].plot(x_plot[:,0],y_plot,'r--',linewidth=2,label='Prediction')
1495                    ax[j].set_title('$t = %.2f$' % (t),fontsize=10)
1496                    ax[j].set_xlabel(' ')
1497                    ax[j].set_ylim([1.3 * ylo.tolist(),1.3 * yup.tolist()])
1498                k = k + 1
1499
1500    #Title
1501    fig.suptitle(title_1d)
1502    fig.tight_layout()
1503
1504    #Show and save
1505    fig = plt.gcf()
1506    if show:
1507        plt.show()
1508    if save:
1509        fig.savefig(file_name + '_slices.png')
1510    plt.close()
1511
1512    #2d plot
1513    if d2:
1514        #Initialize
1515        fig, ax = plt.subplots(1,2)
1516        l1 = jnp.unique(xt[:,-1]).shape[0]
1517        l2 = jnp.unique(xt[:,0]).shape[0]
1518
1519        #Create
1520        ax[0].pcolormesh(xt[:,-1].reshape((l2,l1)),xt[:,0].reshape((l2,l1)),u[:,0].reshape((l2,l1)),cmap = 'RdBu',vmin = ylo.tolist(),vmax = yup.tolist())
1521        ax[0].set_title('Exact')
1522        ax[1].pcolormesh(xt[:,-1].reshape((l2,l1)),xt[:,0].reshape((l2,l1)),upred[:,0].reshape((l2,l1)),cmap = 'RdBu',vmin = ylo.tolist(),vmax = yup.tolist())
1523        ax[1].set_title('Predicted')
1524
1525        #Title
1526        fig.suptitle(title_2d)
1527        fig.tight_layout()
1528
1529        #Show and save
1530        fig = plt.gcf()
1531        if show:
1532            plt.show()
1533        if save:
1534            fig.savefig(file_name + '_2d.png')
1535        plt.close()
1536
1537#Plot results for d = 1
1538def plot_pinn_out2D(times,xt,u,upred,save = False,show = True,file_name = 'result_pinn',title = '',plot_test = True):
1539    """
1540    Plot the prediction of a PINN with 2D output
1541    ----------
1542    Parameters
1543    ----------
1544    times : int
1545
1546        Number of points along the time interval to plot. Default 5
1547
1548    xt : jax.numpy.array
1549
1550        Test data xt array
1551
1552    u : jax.numpy.array
1553
1554        Test data u(x,t) array
1555
1556    upred : jax.numpy.array
1557
1558        Predicted upred(x,t) array on test data
1559
1560    save : logical
1561
1562        Whether to save the plots. Default False
1563
1564    show : logical
1565
1566        Whether to show the plots. Default True
1567
1568    file_name : str
1569
1570        File prefix to save the plots. Default 'result_pinn'
1571
1572    title : str
1573
1574        Title of plot
1575
1576    plot_test : logical
1577
1578        Whether to plot the test data. Default True
1579
1580    Returns
1581    -------
1582    None
1583    """
1584    #Initialize
1585    fig, ax = plt.subplots(int(times/5),5,figsize = (10*int(times/5),3*int(times/5)))
1586    tlo = jnp.min(xt[:,-1])
1587    tup = jnp.max(xt[:,-1])
1588    xlo = jnp.min(u[:,0])
1589    xlo = xlo - 0.1*jnp.abs(xlo)
1590    xup = jnp.max(u[:,0])
1591    xup = xup + 0.1*jnp.abs(xup)
1592    ylo = jnp.min(u[:,1])
1593    ylo = ylo - 0.1*jnp.abs(ylo)
1594    yup = jnp.max(u[:,1])
1595    yup = yup + 0.1*jnp.abs(yup)
1596    k = 0
1597    t_values = np.linspace(tlo,tup,times)
1598
1599    #Create
1600    for i in range(int(times/5)):
1601        for j in range(5):
1602            if k < len(t_values):
1603                t = t_values[k]
1604                t = xt[jnp.abs(xt[:,-1] - t) == jnp.min(jnp.abs(xt[:,-1] - t)),-1][0].tolist()
1605                xpred_plot = upred[xt[:,-1] == t,0]
1606                ypred_plot = upred[xt[:,-1] == t,1]
1607                if plot_test:
1608                    x_plot = u[xt[:,-1] == t,0]
1609                    y_plot = u[xt[:,-1] == t,1]
1610                if int(times/5) > 1:
1611                    if plot_test:
1612                        ax[i,j].plot(x_plot,y_plot,'b-',linewidth=2,label='Exact')
1613                    ax[i,j].plot(xpred_plot,ypred_plot,'r-',linewidth=2,label='Prediction')
1614                    ax[i,j].set_title('$t = %.2f$' % (t),fontsize=10)
1615                    ax[i,j].set_xlabel(' ')
1616                    ax[i,j].set_ylim([1.3 * ylo.tolist(),1.3 * yup.tolist()])
1617                else:
1618                    if plot_test:
1619                        ax[j].plot(x_plot,y_plot,'b-',linewidth=2,label='Exact')
1620                    ax[j].plot(xpred_plot,ypred,'r-',linewidth=2,label='Prediction')
1621                    ax[j].set_title('$t = %.2f$' % (t),fontsize=10)
1622                    ax[j].set_xlabel(' ')
1623                    ax[j].set_ylim([1.3 * ylo.tolist(),1.3 * yup.tolist()])
1624                k = k + 1
1625
1626    #Title
1627    fig.suptitle(title)
1628    fig.tight_layout()
1629
1630    #Show and save
1631    fig = plt.gcf()
1632    if show:
1633        plt.show()
1634    if save:
1635        fig.savefig(file_name + '_slices.png')
1636    plt.close()
1637
1638#Get train data in one array
1639def get_train_data(train_data):
1640    """
1641    Process training sample
1642    ----------
1643
1644    Parameters
1645    ----------
1646    train_data : dict
1647
1648        A dictionay with train data generated by the jinnax.data.generate_PINNdata function
1649
1650    Returns
1651    -------
1652    dict with the processed training data
1653    """
1654    xdata = None
1655    ydata = None
1656    xydata = None
1657    if train_data['sensor'] is not None:
1658        sensor_sample = train_data['sensor'].shape[0]
1659        xdata = train_data['sensor']
1660        ydata = train_data['usensor']
1661        xydata = jnp.column_stack((train_data['sensor'],train_data['usensor']))
1662    else:
1663        sensor_sample = 0
1664    if train_data['boundary'] is not None:
1665        boundary_sample = train_data['boundary'].shape[0]
1666        if xdata is not None:
1667            xdata = jnp.vstack((xdata,train_data['boundary']))
1668            ydata = jnp.vstack((ydata,train_data['uboundary']))
1669            xydata = jnp.vstack((xydata,jnp.column_stack((train_data['boundary'],train_data['uboundary']))))
1670        else:
1671            xdata = train_data['boundary']
1672            ydata = train_data['uboundary']
1673            xydata = jnp.column_stack((train_data['boundary'],train_data['uboundary']))
1674    else:
1675        boundary_sample = 0
1676    if train_data['initial'] is not None:
1677        initial_sample = train_data['initial'].shape[0]
1678        if xdata is not None:
1679            xdata = jnp.vstack((xdata,train_data['initial']))
1680            ydata = jnp.vstack((ydata,train_data['uinitial']))
1681            xydata = jnp.vstack((xydata,jnp.column_stack((train_data['initial'],train_data['uinitial']))))
1682        else:
1683            xdata = train_data['initial']
1684            ydata = train_data['uinitial']
1685            xydata = jnp.column_stack((train_data['initial'],train_data['uinitial']))
1686    else:
1687        initial_sample = 0
1688    if train_data['collocation'] is not None:
1689        collocation_sample = train_data['collocation'].shape[0]
1690    else:
1691        collocation_sample = 0
1692
1693    return {'xy': xydata,'x': xdata,'y': ydata,'sensor_sample': sensor_sample,'boundary_sample': boundary_sample,'initial_sample': initial_sample,'collocation_sample': collocation_sample}
1694
1695#Process training
1696def process_training(test_data,file_name,at_each = 100,bolstering = True,mc_sample = 10000,save = False,file_name_save = 'result_pinn',key = 0,ec = 1e-6,lamb = 1):
1697    """
1698    Process the training of a Physics-informed Neural Network
1699    ----------
1700
1701    Parameters
1702    ----------
1703    test_data : dict
1704
1705        A dictionay with test data for L2 error calculation generated by the jinnax.data.generate_PINNdata function
1706
1707    file_name : str
1708
1709        Name of the files saved during training
1710
1711    at_each : int
1712
1713        Compute results for epochs multiple of at_each. Default 100
1714
1715    bolstering : logical
1716
1717        Whether to compute bolstering mean square error. Default True
1718
1719    mc_sample : int
1720
1721        Number of sample for Monte Carlo integration in bolstering. Default 10000
1722
1723    save : logical
1724
1725        Whether to save the training results. Default False
1726
1727    file_name_save : str
1728
1729        File prefix to save the plots and the L2 error. Default 'result_pinn'
1730
1731    key : int
1732
1733        Key for random samples in bolstering. Default 0
1734
1735    ec : float
1736
1737        Stopping criteria error for EM algorithm in bolstering. Default 1e-6
1738
1739    lamb : float
1740
1741        Hyperparameter of EM algorithm in bolstering. Default 1
1742
1743    Returns
1744    -------
1745    pandas data frame with training results
1746    """
1747    #Config
1748    config = pickle.load(open(file_name + '_config.pickle', 'rb'))
1749    epochs = config['epochs']
1750    train_data = config['train_data']
1751    forward = config['forward']
1752
1753    #Get train data
1754    td = get_train_data(train_data)
1755    xydata = td['xy']
1756    xdata = td['x']
1757    ydata = td['y']
1758    sensor_sample = td['sensor_sample']
1759    boundary_sample = td['boundary_sample']
1760    initial_sample = td['initial_sample']
1761    collocation_sample = td['collocation_sample']
1762
1763    #Generate keys
1764    if bolstering:
1765        keys = jax.random.split(jax.random.PRNGKey(key),epochs)
1766
1767    #Initialize loss
1768    train_mse = []
1769    test_mse = []
1770    train_L2 = []
1771    test_L2 = []
1772    bolstX = []
1773    bolstXY = []
1774    loss = []
1775    time = []
1776    ep = []
1777
1778    #Process training
1779    with alive_bar(epochs) as bar:
1780        for e in range(epochs):
1781            if (e % at_each == 0 and at_each != epochs) or e == epochs - 1:
1782                ep = ep + [e]
1783
1784                #Read parameters
1785                params = pickle.load(open(file_name + '_epoch' + str(e).rjust(6, '0') + '.pickle','rb'))
1786
1787                #Time
1788                time = time + [params['time']]
1789
1790                #Define learned function
1791                def psi(x):
1792                    return forward(x,params['params']['net'])
1793
1794                #Train MSE and L2
1795                if xdata is not None:
1796                    train_mse = train_mse + [jnp.mean(MSE(psi(xdata),ydata)).tolist()]
1797                    train_L2 = train_L2 + [L2error(psi(xdata),ydata).tolist()]
1798                else:
1799                    train_mse = train_mse + [None]
1800                    train_L2 = train_L2 + [None]
1801
1802                #Test MSE and L2
1803                test_mse = test_mse + [jnp.mean(MSE(psi(test_data['xt']),test_data['u'])).tolist()]
1804                test_L2 = test_L2 + [L2error(psi(test_data['xt']),test_data['u']).tolist()]
1805
1806                #Bolstering
1807                if bolstering:
1808                    bX = []
1809                    bXY = []
1810                    for method in ['chi','mm','mpe']:
1811                        kxy = gk.kernel_estimator(data = xydata,key = keys[e,0],method = method,lamb = lamb,ec = ec,psi = psi)
1812                        kx = gk.kernel_estimator(data = xdata,key = keys[e,0],method = method,lamb = lamb,ec = ec,psi = psi)
1813                        bX = bX + [gb.bolstering(psi,xdata,ydata,kx,key = keys[e,0],mc_sample = mc_sample).tolist()]
1814                        bXY = bXY + [gb.bolstering(psi,xdata,ydata,kxy,key = keys[e,0],mc_sample = mc_sample).tolist()]
1815                    for bias in [1/jnp.sqrt(xdata.shape[0]),1/xdata.shape[0],1/(xdata.shape[0] ** 2),1/(xdata.shape[0] ** 3),1/(xdata.shape[0] ** 4)]:
1816                        kx = gk.kernel_estimator(data = xydata,key = keys[e,0],method = 'hessian',lamb = lamb,ec = ec,psi = psi,bias = bias)
1817                        bX = bX + [gb.bolstering(psi,xdata,ydata,kx,key = keys[e,0],mc_sample = mc_sample).tolist()]
1818                    bolstX = bolstX + [bX]
1819                    bolstXY = bolstXY + [bXY]
1820                else:
1821                    bolstX = bolstX + [None]
1822                    bolstXY = bolstXY + [None]
1823
1824                #Loss
1825                loss = loss + [params['loss'].tolist()]
1826
1827                #Delete
1828                del params, psi
1829            #Update alive_bar
1830            bar()
1831
1832    #Bolstering results
1833    if bolstering:
1834        bolstX = jnp.array(bolstX)
1835        bolstXY = jnp.array(bolstXY)
1836
1837    #Create data frame
1838    if bolstering:
1839        df = pd.DataFrame(np.column_stack([ep,time,[sensor_sample] * len(ep),[boundary_sample] * len(ep),[initial_sample] * len(ep),[collocation_sample] * len(ep),loss,
1840            train_mse,test_mse,train_L2,test_L2,bolstX[:,0],bolstXY[:,0],bolstX[:,1],bolstXY[:,1],bolstX[:,2],bolstXY[:,2],bolstX[:,3],bolstX[:,4],bolstX[:,5],bolstX[:,6],bolstX[:,7]]),
1841            columns=['epoch','training_time','sensor_sample','boundary_sample','initial_sample','collocation_sample','loss','train_mse','test_mse','train_L2','test_L2','bolstX_chi','bolstXY_chi','bolstX_mm','bolstXY_mm','bolstX_mpe','bolstXY_mpe','bolstHessian_sqrtn','bolstHessian_n','bolstHessian_n2','bolstHessian_n3','bolstHessian_n4'])
1842    else:
1843        df = pd.DataFrame(np.column_stack([ep,time,[sensor_sample] * len(ep),[boundary_sample] * len(ep),[initial_sample] * len(ep),[collocation_sample] * len(ep),loss,
1844            train_mse,test_mse,train_L2,test_L2]),
1845            columns=['epoch','training_time','sensor_sample','boundary_sample','initial_sample','collocation_sample','loss','train_mse','test_mse','train_L2','test_L2'])
1846    if save:
1847        df.to_csv(file_name_save + '.csv',index = False)
1848
1849    return df
1850
1851#Demo video for training1D PINN
1852def demo_train_pinn1D(test_data,file_name,at_each = 100,times = 5,d2 = True,file_name_save = 'result_pinn_demo',title = '',framerate = 10):
1853    """
1854    Demo video with the training of a 1D PINN
1855    ----------
1856
1857    Parameters
1858    ----------
1859    test_data : dict
1860
1861        A dictionay with test data for L2 error calculation generated by the jinnax.data.generate_PINNdata function
1862
1863    file_name : str
1864
1865        Name of the files saved during training
1866
1867    at_each : int
1868
1869        Compute results for epochs multiple of at_each. Default 100
1870
1871    times : int
1872
1873        Number of points along the time interval to plot. Default 5
1874
1875    d2 : logical
1876
1877        Whether to make video demo of 2D plot. Default True
1878
1879    file_name_save : str
1880
1881        File prefix to save the plots and videos. Default 'result_pinn_demo'
1882
1883    title : str
1884
1885        Title for plots
1886
1887    framerate : int
1888
1889        Framerate for video. Default 10
1890
1891    Returns
1892    -------
1893    None
1894    """
1895    #Config
1896    with open(file_name + '_config.pickle', 'rb') as file:
1897        config = pickle.load(file)
1898    epochs = config['epochs']
1899    train_data = config['train_data']
1900    forward = config['forward']
1901
1902    #Get train data
1903    td = get_train_data(train_data)
1904    xt = td['x']
1905    u = td['y']
1906
1907    #Create folder to save plots
1908    os.system('mkdir ' + file_name_save)
1909
1910    #Create images
1911    k = 1
1912    with alive_bar(epochs) as bar:
1913        for e in range(epochs):
1914            if e % at_each == 0 or e == epochs - 1:
1915                #Read parameters
1916                params = pd.read_pickle(file_name + '_epoch' + str(e).rjust(6, '0') + '.pickle')
1917
1918                #Define learned function
1919                def psi(x):
1920                    return forward(x,params['params']['net'])
1921
1922                #Compute L2 train, L2 test and loss
1923                loss = params['loss']
1924                L2_train = L2error(psi(xt),u)
1925                L2_test = L2error(psi(test_data['xt']),test_data['u'])
1926                title_epoch = title + ' Epoch = ' + str(e) + ' L2 train = ' + str(round(L2_train,6)) + ' L2 test = ' + str(round(L2_test,6))
1927
1928                #Save plot
1929                plot_pinn1D(times,test_data['xt'],test_data['u'],psi(test_data['xt']),d2,save = True,show = False,file_name = file_name_save + '/' + str(k),title_1d = title_epoch,title_2d = title_epoch)
1930                k = k + 1
1931
1932                #Delete
1933                del params, psi, loss, L2_train, L2_test, title_epoch
1934            #Update alive_bar
1935            bar()
1936    #Create demo video
1937    os.system('ffmpeg -framerate ' + str(framerate) + ' -i ' + file_name_save + '/' + '%00d_slices.png -c:v libx264 -profile:v high -crf 20 -pix_fmt yuv420p ' + file_name_save + '/' + file_name_save + '_slices.mp4')
1938    if d2:
1939        os.system('ffmpeg -framerate ' + str(framerate) + ' -i ' + file_name_save + '/' + '%00d_2d.png -c:v libx264 -profile:v high -crf 20 -pix_fmt yuv420p ' + file_name_save + '/' + file_name_save + '_2d.mp4')
1940
1941#Demo in time for 1D PINN
1942def demo_time_pinn1D(test_data,file_name,epochs,file_name_save = 'result_pinn_time_demo',title = '',framerate = 10):
1943    """
1944    Demo video with the time evolution of a 1D PINN
1945    ----------
1946
1947    Parameters
1948    ----------
1949    test_data : dict
1950
1951        A dictionay with test data for L2 error calculation generated by the jinnax.data.generate_PINNdata function
1952
1953    file_name : str
1954
1955        Name of the files saved during training
1956
1957    epochs : list
1958
1959        Which training epochs to plot
1960
1961    file_name_save : str
1962
1963        File prefix to save the plots and video. Default 'result_pinn_time_demo'
1964
1965    title : str
1966
1967        Title for plots
1968
1969    framerate : int
1970
1971        Framerate for video. Default 10
1972
1973    Returns
1974    -------
1975    None
1976    """
1977    #Config
1978    with open(file_name + '_config.pickle', 'rb') as file:
1979        config = pickle.load(file)
1980    train_data = config['train_data']
1981    forward = config['forward']
1982
1983    #Create folder to save plots
1984    os.system('mkdir ' + file_name_save)
1985
1986    #Plot parameters
1987    tdom = jnp.unique(test_data['xt'][:,-1])
1988    ylo = jnp.min(test_data['u'])
1989    ylo = ylo - 0.1*jnp.abs(ylo)
1990    yup = jnp.max(test_data['u'])
1991    yup = yup + 0.1*jnp.abs(yup)
1992
1993    #Open PINN for each epoch
1994    results = []
1995    upred = []
1996    for e in epochs:
1997        tmp = pd.read_pickle(file_name + '_epoch' + str(e).rjust(6, '0') + '.pickle')
1998        results = results + [tmp]
1999        upred = upred + [forward(test_data['xt'],tmp['params']['net'])]
2000
2001    #Create images
2002    k = 1
2003    with alive_bar(len(tdom)) as bar:
2004        for t in tdom:
2005            #Test data
2006            xt_step = test_data['xt'][test_data['xt'][:,-1] == t]
2007            u_step = test_data['u'][test_data['xt'][:,-1] == t]
2008            #Initialize plot
2009            if len(epochs) > 1:
2010                fig, ax = plt.subplots(int(len(epochs)/2),2,figsize = (10,5*len(epochs)/2))
2011            else:
2012                fig, ax = plt.subplots(1,1,figsize = (10,5))
2013            #Create
2014            index = 0
2015            if int(len(epochs)/2) > 1:
2016                for i in range(int(len(epochs)/2)):
2017                    for j in range(min(2,len(epochs))):
2018                        upred_step = upred[index][test_data['xt'][:,-1] == t]
2019                        ax[i,j].plot(xt_step[:,0],u_step[:,0],'b-',linewidth=2,label='Exact')
2020                        ax[i,j].plot(xt_step[:,0],upred_step[:,0],'r--',linewidth=2,label='Prediction')
2021                        ax[i,j].set_title('Epoch = ' + str(epochs[index]),fontsize=10)
2022                        ax[i,j].set_xlabel(' ')
2023                        ax[i,j].set_ylim([1.3 * ylo.tolist(),1.3 * yup.tolist()])
2024                        index = index + 1
2025            elif len(epochs) > 1:
2026                for j in range(2):
2027                    upred_step = upred[index][test_data['xt'][:,-1] == t]
2028                    ax[j].plot(xt_step[:,0],u_step[:,0],'b-',linewidth=2,label='Exact')
2029                    ax[j].plot(xt_step[:,0],upred_step[:,0],'r--',linewidth=2,label='Prediction')
2030                    ax[j].set_title('Epoch = ' + str(epochs[index]),fontsize=10)
2031                    ax[j].set_xlabel(' ')
2032                    ax[j].set_ylim([1.3 * ylo.tolist(),1.3 * yup.tolist()])
2033                    index = index + 1
2034            else:
2035                upred_step = upred[index][test_data['xt'][:,-1] == t]
2036                ax.plot(xt_step[:,0],u_step[:,0],'b-',linewidth=2,label='Exact')
2037                ax.plot(xt_step[:,0],upred_step[:,0],'r--',linewidth=2,label='Prediction')
2038                ax.set_title('Epoch = ' + str(epochs[index]),fontsize=10)
2039                ax.set_xlabel(' ')
2040                ax.set_ylim([1.3 * ylo.tolist(),1.3 * yup.tolist()])
2041                index = index + 1
2042
2043
2044            #Title
2045            fig.suptitle(title + 't = ' + str(round(t,4)))
2046            fig.tight_layout()
2047
2048            #Show and save
2049            fig = plt.gcf()
2050            fig.savefig(file_name_save + '/' + str(k) + '.png')
2051            k = k + 1
2052            plt.close()
2053            bar()
2054
2055    #Create demo video
2056    os.system('ffmpeg -framerate ' + str(framerate) + ' -i ' + file_name_save + '/' + '%00d.png -c:v libx264 -profile:v high -crf 20 -pix_fmt yuv420p ' + file_name_save + '/' + file_name_save + '_time_demo.mp4')