fennol.models.misc.encodings

  1import jax
  2import jax.numpy as jnp
  3import flax.linen as nn
  4from typing import Optional, Union, List, Sequence, ClassVar, Dict
  5import math
  6import dataclasses
  7import numpy as np
  8from .nets import FullyConnectedNet
  9from ...utils import AtomicUnits as au
 10from functools import partial
 11from ...utils.periodic_table import (
 12    PERIODIC_TABLE_REV_IDX,
 13    PERIODIC_TABLE,
 14    EL_STRUCT,
 15    VALENCE_STRUCTURE,
 16    XENONPY_PROPS,
 17    SJS_COORDINATES,
 18    PERIODIC_COORDINATES,
 19    ATOMIC_IONIZATION_ENERGY,
 20    POLARIZABILITIES,
 21    D3_COV_RADII,
 22    VDW_RADII,
 23    OXIDATION_STATES,
 24)
 25
 26
 27class SpeciesEncoding(nn.Module):
 28    """A module that encodes chemical species information.
 29
 30    FID: SPECIES_ENCODING
 31    """
 32
 33    encoding: str = "random"
 34    """ The encoding to use. Can be one of "one_hot", "occupation", "electronic_structure", "properties", "sjs_coordinates", "random". 
 35        Multiple encodings can be concatenated using the "+" separator.
 36    """
 37    dim: int = 16
 38    """ The dimension of the encoding if not fixed by design."""
 39    zmax: int = 86
 40    """ The maximum atomic number to encode."""
 41    output_key: Optional[str] = None
 42    """ The key to use for the output in the returned dictionary."""
 43
 44    species_order: Optional[Union[str, Sequence[str]]] = None
 45    """ The order of the species to use for the encoding. Only used for "onehot" encoding.
 46         If None, we encode all elements up to `zmax`."""
 47    trainable: bool = False
 48    """ Whether the encoding is trainable or fixed. Does not apply to "random" encoding which is always trainable."""
 49    extra_params: Dict = dataclasses.field(default_factory=dict)
 50    """ Dictionary of extra parameters for the basis."""
 51
 52    FID: ClassVar[str] = "SPECIES_ENCODING"
 53
 54    @nn.compact
 55    def __call__(self, inputs: Union[dict, jax.Array]) -> Union[dict, jax.Array]:
 56
 57        zmax = self.zmax
 58        if zmax <= 0 or zmax > len(PERIODIC_TABLE):
 59            zmax = len(PERIODIC_TABLE)
 60
 61        zmaxpad = zmax + 2
 62
 63        encoding = self.encoding.lower().strip()
 64        encodings = encoding.split("+")
 65        ############################
 66        conv_tensors = []
 67
 68        if "one_hot" in encodings or "onehot" in encodings:
 69            if self.species_order is None:
 70                conv_tensor = np.eye(zmax)
 71                conv_tensor = np.concatenate(
 72                    [np.zeros((1, zmax)), conv_tensor, np.zeros((1, zmax))], axis=0
 73                )
 74            else:
 75                if isinstance(self.species_order, str):
 76                    species_order = [el.strip() for el in self.species_order.split(",")]
 77                else:
 78                    species_order = [el for el in self.species_order]
 79                conv_tensor = np.zeros((zmaxpad, len(species_order)))
 80                for i, s in enumerate(species_order):
 81                    conv_tensor[PERIODIC_TABLE_REV_IDX[s], i] = 1
 82
 83            conv_tensors.append(conv_tensor)
 84
 85        if "occupation" in encodings:
 86            conv_tensor = np.zeros((zmaxpad, (zmax + 1) // 2))
 87            for i in range(1, zmax + 1):
 88                conv_tensor[i, : i // 2] = 1
 89                if i % 2 == 1:
 90                    conv_tensor[i, i // 2] = 0.5
 91
 92            conv_tensors.append(conv_tensor)
 93
 94        if "electronic_structure" in encodings:
 95            Z = np.arange(1, zmax + 1).reshape(-1, 1)
 96            Zref = [zmax]
 97            e_struct = np.array(EL_STRUCT[1 : zmax + 1])
 98            eref = [2, 2, 6, 2, 6, 2, 10, 6, 2, 10, 6, 2, 14, 10, 6, 2, 14, 10, 6]
 99            v_struct = np.array(VALENCE_STRUCTURE[1 : zmax + 1])
100            vref = [2, 6, 10, 14]
101            if zmax <= 86:
102                e_struct = e_struct[:, :15]
103                eref = eref[:15]
104            ref = np.array(Zref + eref + vref)
105            conv_tensor = np.concatenate([Z, e_struct, v_struct], axis=1)
106            conv_tensor = conv_tensor / ref[None, :]
107            dim = conv_tensor.shape[1]
108            conv_tensor = np.concatenate(
109                [np.zeros((1, dim)), conv_tensor, np.zeros((1, dim))], axis=0
110            )
111
112            conv_tensors.append(conv_tensor)
113
114        if "properties" in encodings:
115            props = np.array(XENONPY_PROPS)[1:-1]
116            assert (
117                self.zmax <= props.shape[0]
118            ), f"zmax > {props.shape[0]} not supported for xenonpy properties"
119            conv_tensor = props[1 : zmax + 1]
120            mean = np.mean(props, axis=0)
121            std = np.std(props, axis=0)
122            conv_tensor = (conv_tensor - mean[None, :]) / std[None, :]
123            dim = conv_tensor.shape[1]
124            conv_tensor = np.concatenate(
125                [np.zeros((1, dim)), conv_tensor, np.zeros((1, dim))], axis=0
126            )
127            conv_tensors.append(conv_tensor)
128
129        if "valence_properties" in encodings:
130            assert zmax <= 86, "Valence properties only available for zmax <= 86"
131            Z = np.arange(1, zmax + 1).reshape(-1, 1)
132            Zref = [zmax]
133            Zinv = 1.0 / Z
134            Zinvref = [1.0]
135            v_struct = np.array(VALENCE_STRUCTURE[1 : zmax + 1])
136            vref = [2, 6, 10, 14]
137            ionization = np.array(ATOMIC_IONIZATION_ENERGY[1 : zmax + 1]).reshape(-1, 1)
138            ionizationref = [0.5]
139            polariz = np.array(POLARIZABILITIES[1 : zmax + 1]).reshape(-1, 1)
140            polarizref = [np.median(polariz)]
141
142            cov = np.array(D3_COV_RADII[1 : zmax + 1]).reshape(-1, 1)
143            covref = [np.median(cov)]
144
145            vdw = np.array(VDW_RADII[1 : zmax + 1]).reshape(-1, 1)
146            vdwref = [np.median(vdw)]
147
148            ref = np.array(
149                Zref + Zinvref + ionizationref + polarizref + covref + vdwref + vref
150            )
151            conv_tensor = np.concatenate(
152                [Z, Zinv, ionization, polariz, cov, vdw, v_struct], axis=1
153            )
154
155            conv_tensor = conv_tensor / ref[None, :]
156            dim = conv_tensor.shape[1]
157            conv_tensor = np.concatenate(
158                [np.zeros((1, dim)), conv_tensor, np.zeros((1, dim))], axis=0
159            )
160            conv_tensors.append(conv_tensor)
161
162        if "sjs_coordinates" in encodings:
163            coords = np.array(SJS_COORDINATES)[1:-1]
164            conv_tensor = coords[1 : zmax + 1]
165            mean = np.mean(coords, axis=0)
166            std = np.std(coords, axis=0)
167            conv_tensor = (conv_tensor - mean[None, :]) / std[None, :]
168            dim = conv_tensor.shape[1]
169            conv_tensor = np.concatenate(
170                [np.zeros((1, dim)), conv_tensor, np.zeros((1, dim))], axis=0
171            )
172            conv_tensors.append(conv_tensor)
173
174        if "positional" in encodings:
175            coords = np.array(PERIODIC_COORDINATES)[1 : zmax + 1]
176            row, col = coords[:, 0], coords[:, 1]
177            drow = self.extra_params.get("drow", default=5)
178            dcol = self.dim - drow
179            nrow = self.extra_params.get("nrow", default=100.0)
180            ncol = self.extra_params.get("ncol", default=1000.0)
181
182            erow = positional_encoding_static(row, drow, nrow)
183            ecol = positional_encoding_static(col, dcol, ncol)
184            conv_tensor = np.concatenate([erow, ecol], axis=-1)
185            dim = conv_tensor.shape[1]
186            conv_tensor = np.concatenate(
187                [np.zeros((1, dim)), conv_tensor, np.zeros((1, dim))], axis=0
188            )
189            conv_tensors.append(conv_tensor)
190
191        if "oxidation" in encodings:
192            states_set = sorted(set(sum(OXIDATION_STATES, [])))
193            nstates = len(states_set)
194            state_dict = {s: i for i, s in enumerate(states_set)}
195            conv_tensor = np.zeros((zmaxpad, nstates))
196            for i, states in enumerate(OXIDATION_STATES[1 : zmax + 1]):
197                for s in states:
198                    conv_tensor[i, state_dict[s]] = 1
199            conv_tensors.append(conv_tensor)
200
201        if len(conv_tensors) > 0:
202            conv_tensor = np.concatenate(conv_tensors, axis=1)
203            if self.trainable:
204                conv_tensor = self.param(
205                    "conv_tensor",
206                    lambda key: jnp.asarray(conv_tensor, dtype=jnp.float32),
207                )
208            else:
209                conv_tensor = jnp.asarray(conv_tensor, dtype=jnp.float32)
210            conv_tensors = [conv_tensor]
211        else:
212            conv_tensors = []
213
214        if "random" in encodings:
215            rand_encoding = self.param(
216                "rand_encoding",
217                lambda key, shape: jax.nn.standardize(
218                    jax.random.normal(key, shape, dtype=jnp.float32)
219                ),
220                (zmaxpad, self.dim),
221            )
222            conv_tensors.append(rand_encoding)
223
224        if "randint" in encodings:
225            rand_encoding = self.param(
226                "randint_encoding",
227                lambda key, shape: jax.random.randint(key, shape, 0, 2).astype(
228                    jnp.float32
229                ),
230                (zmaxpad, self.dim),
231            )
232            conv_tensors.append(rand_encoding)
233
234        if "randtri" in encodings:
235            rand_encoding = self.param(
236                "randtri_encoding",
237                lambda key, shape: (jax.random.randint(key, shape, 0, 3) - 1).astype(
238                    jnp.float32
239                ),
240                (zmaxpad, self.dim),
241            )
242            conv_tensors.append(rand_encoding)
243
244        assert len(conv_tensors) > 0, f"No encoding recognized in '{self.encoding}'"
245
246        conv_tensor = jnp.concatenate(conv_tensors, axis=1)
247
248        species = inputs["species"] if isinstance(inputs, dict) else inputs
249        out = conv_tensor[species]
250        ############################
251
252        if isinstance(inputs, dict):
253            output_key = self.name if self.output_key is None else self.output_key
254            out = out.astype(inputs["coordinates"].dtype)
255            return {**inputs, output_key: out} if output_key is not None else out
256        return out
257
258
259class RadialBasis(nn.Module):
260    """Computes a radial encoding of distances.
261
262    FID: RADIAL_BASIS
263    """
264
265    end: float
266    """ The maximum distance to consider."""
267    start: float = 0.0
268    """ The minimum distance to consider."""
269    dim: int = 8
270    """ The dimension of the basis."""
271    graph_key: Optional[str] = None
272    """ The key of the graph in the inputs."""
273    output_key: Optional[str] = None
274    """ The key to use for the output in the returned dictionary."""
275    basis: str = "bessel"
276    """ The basis to use. Can be one of "bessel", "gaussian", "gaussian_rinv", "fourier", "spooky"."""
277    trainable: bool = False
278    """ Whether the basis parameters are trainable or fixed."""
279    enforce_positive: bool = False
280    """ Whether to enforce distance-start to be positive"""
281    gamma: float = 1.0 / (2 * au.BOHR)
282    """ The gamma parameter for the "spooky" basis."""
283    n_levels: int = 10
284    """ The number of levels for the "levels" basis."""
285    alt_bessel_norm: bool = False
286    """ If True, use the (2/(end-start))**0.5 normalization for the bessel basis."""
287    extra_params: Dict = dataclasses.field(default_factory=dict)
288    """ Dictionary of extra parameters for the basis."""
289
290    FID: ClassVar[str] = "RADIAL_BASIS"
291
292    @nn.compact
293    def __call__(self, inputs: Union[dict, jax.Array]) -> Union[dict, jax.Array]:
294        if self.graph_key is not None:
295            x = inputs[self.graph_key]["distances"]
296        else:
297            x = inputs["distances"] if isinstance(inputs, dict) else inputs
298
299        shape = x.shape
300        x = x.reshape(-1)
301
302        basis = self.basis.lower()
303        ############################
304        if basis == "bessel":
305            c = self.end - self.start
306            x = x[:, None] - self.start
307            # if self.enforce_positive:
308            #     x = jax.nn.softplus(x)
309
310            if self.trainable:
311                bessel_roots = self.param(
312                    "bessel_roots",
313                    lambda key, dim: jnp.asarray(
314                        np.arange(1, dim + 1, dtype=x.dtype)[None, :] * (math.pi / c)
315                    ),
316                    self.dim,
317                )
318                norm = 1.0 / jnp.max(
319                    bessel_roots
320                )  # (2.0 / c) ** 0.5 /jnp.max(bessel_roots)
321            else:
322                bessel_roots = jnp.asarray(
323                    np.arange(1, self.dim + 1, dtype=x.dtype)[None, :] * (math.pi / c)
324                )
325                norm = 1.0 / (
326                    self.dim * math.pi / c
327                )  # (2.0 / c) ** 0.5/(self.dim*math.pi/c)
328            if self.alt_bessel_norm:
329                norm = (2.0 / c) ** 0.5
330            out = norm * jnp.sin(x * bessel_roots) / x
331
332            if self.enforce_positive:
333                out = jnp.where(x > 0, out * (1.0 - jnp.exp(-(x**2))), 0.0)
334
335        elif basis == "gaussian":
336            if self.trainable:
337                roots = self.param(
338                    "radial_centers",
339                    lambda key, dim, start, end: jnp.linspace(
340                        start, end, dim + 1, dtype=x.dtype
341                    )[None, :-1],
342                    self.dim,
343                    self.start,
344                    self.end,
345                )
346                eta = self.param(
347                    "radial_etas",
348                    lambda key, dim, start, end: jnp.full(
349                        dim,
350                        dim / (end - start),
351                        dtype=x.dtype,
352                    )[None, :],
353                    self.dim,
354                    self.start,
355                    self.end,
356                )
357
358            else:
359                roots = jnp.asarray(
360                    np.linspace(self.start, self.end, self.dim + 1)[None, :-1],
361                    dtype=x.dtype,
362                )
363                eta = jnp.asarray(
364                    np.full(self.dim, self.dim / (self.end - self.start))[None, :],
365                    dtype=x.dtype,
366                )
367
368            x = x[:, None]
369            x2 = (eta * (x - roots)) ** 2
370            out = jnp.exp(-x2)
371            if self.enforce_positive:
372                out = jnp.where(
373                    x > self.start,
374                    out * (1.0 - jnp.exp(-10 * (x - self.start) ** 2)),
375                    0.0,
376                )
377
378        elif basis == "gaussian_rinv":
379            rinv_high = 1.0 / max(self.start, 0.1)
380            rinv_low = 1.0 / (0.8 * self.end)
381
382            if self.trainable:
383                roots = self.param(
384                    "radial_centers",
385                    lambda key, dim, rinv_low, rinv_high: jnp.linspace(
386                        rinv_low, rinv_high, dim, dtype=x.dtype
387                    )[None, :],
388                    self.dim,
389                    rinv_low,
390                    rinv_high,
391                )
392                sigmas = self.param(
393                    "radial_sigmas",
394                    lambda key, dim, rinv_low, rinv_high: jnp.full(
395                        dim,
396                        2**0.5 / (2 * dim * rinv_low),
397                        dtype=x.dtype,
398                    )[None, :],
399                    self.dim,
400                    rinv_low,
401                    rinv_high,
402                )
403            else:
404                roots = jnp.asarray(
405                    np.linspace(rinv_low, rinv_high, self.dim, dtype=x.dtype)[None, :]
406                )
407                sigmas = jnp.asarray(
408                    np.full(
409                        self.dim,
410                        2**0.5 / (2 * self.dim * rinv_low),
411                    )[None, :],
412                    dtype=x.dtype,
413                )
414
415            rinv = 1.0 / x[:, None]
416
417            out = jnp.exp(-((rinv - roots) ** 2) / sigmas**2)
418
419        elif basis == "fourier":
420            if self.trainable:
421                roots = self.param(
422                    "roots",
423                    lambda key, dim: jnp.arange(dim, dtype=x.dtype)[None, :] * math.pi,
424                    self.dim,
425                )
426            else:
427                roots = jnp.asarray(
428                    np.arange(self.dim)[None, :] * math.pi, dtype=x.dtype
429                )
430            c = self.end - self.start
431            x = x[:, None] - self.start
432            # if self.enforce_positive:
433            #     x = jax.nn.softplus(x)
434            norm = 1.0 / (0.25 + 0.5 * self.dim) ** 0.5
435            out = norm * jnp.cos(x * roots / c)
436            if self.enforce_positive:
437                out = jnp.where(x > 0, out, norm)
438
439        elif basis == "spooky":
440
441            gamma = self.gamma
442            if self.trainable:
443                gamma = jnp.abs(
444                    self.param("gamma", lambda key: jnp.asarray(gamma, dtype=x.dtype))
445                )
446
447            if self.enforce_positive:
448                x = jnp.where(x - self.start > 1.0e-3, x - self.start, 1.0e-3)[:, None]
449                dim = self.dim
450            else:
451                x = x[:, None] - self.start
452                dim = self.dim - 1
453
454            norms = []
455            for k in range(self.dim):
456                norms.append(math.comb(dim, k))
457            norms = jnp.asarray(np.array(norms)[None, :], dtype=x.dtype)
458
459            e = jnp.exp(-gamma * x)
460            k = jnp.asarray(np.arange(self.dim, dtype=x.dtype)[None, :])
461            b = e**k * (1 - e) ** (dim - k)
462            out = b * e * norms
463            if self.enforce_positive:
464                out = jnp.where(x > 1.0e-3, out * (1.0 - jnp.exp(-(x**2))), 0.0)
465            # logfac = np.zeros(self.dim)
466            # for i in range(2, self.dim):
467            #     logfac[i] = logfac[i - 1] + np.log(i)
468            # k = np.arange(self.dim)
469            # n = self.dim - 1 - k
470            # logbin = jnp.asarray((logfac[-1] - logfac[k] - logfac[n])[None,:], dtype=x.dtype)
471            # n = jnp.asarray(n[None,:], dtype=x.dtype)
472            # k = jnp.asarray(k[None,:], dtype=x.dtype)
473
474            # gamma = 1.0 / (2 * au.BOHR)
475            # if self.trainable:
476            #     gamma = jnp.abs(
477            #         self.param("gamma", lambda key: jnp.asarray(gamma, dtype=x.dtype))
478            #     )
479            # gammar = (-gamma * x)[:,None]
480            # x = logbin + n * gammar + k * jnp.log(-jnp.expm1(gammar))
481            # out = jnp.exp(x)*jnp.exp(gammar)
482        elif basis == "levels":
483            assert self.n_levels >= 2, "Number of levels must be >= 2."
484
485            def initialize_levels(key):
486                key0, key1, key_phi = jax.random.split(key, 3)
487                level0 = jax.random.randint(key0, (self.dim,), 0, 2)
488                level1 = jax.random.randint(key1, (self.dim,), 0, 2)
489                # level0 = jax.random.normal(key0, (self.dim,), dtype=jnp.float32)
490                # level1 = jax.random.normal(key1, (self.dim,), dtype=jnp.float32)
491                phi = jax.random.uniform(key_phi, (self.dim,), dtype=jnp.float32)
492                levels = [level0]
493                for l in range(2, self.n_levels - 1):
494                    tau = float(self.n_levels - l) / float(self.n_levels - 1)
495                    phil = phi < tau
496                    level = jnp.where(phil, level0, level1)
497                    levels.append(level)
498                levels.append(level1)
499                return jnp.stack(levels).astype(jnp.float32)
500
501            levels = self.param("levels", initialize_levels)
502            # levels = self.param(
503            #     "levels",
504            #     lambda key, shape: jax.random.normal(key, shape, dtype=jnp.float32),
505            #     (self.n_levels,self.dim),
506            # )
507
508            flevel = (x - self.start) / (self.end - self.start) * (self.n_levels - 1)
509            ilevel = jnp.floor(flevel).astype(jnp.int32)
510            ilevel1 = jnp.clip(ilevel + 1, 0, self.n_levels - 1)
511            ilevel = jnp.clip(ilevel, 0, self.n_levels - 1)
512
513            dx = flevel - ilevel
514            w = 0.5 * (1 + jnp.cos(jnp.pi * dx))[:, None]
515
516            ## interpolate between level vectors
517            v1 = levels[ilevel]
518            v2 = levels[ilevel1]
519            out = v1 * w + v2 * (1 - w)
520        elif basis == "finite_support":
521            flevel = (x - self.start) / (self.end - self.start) * (self.dim + 1)
522            ilevel = jnp.floor(flevel).astype(jnp.int32)
523            ilevel1 = jnp.clip(ilevel + 1, 0, self.dim + 1)
524            ilevel = jnp.clip(ilevel, 0, self.dim + 1)
525
526            dx = flevel - ilevel
527            w = 0.5 * (1 + jnp.cos(jnp.pi * dx))
528
529            ilevelflat = ilevel + jnp.arange(x.shape[0]) * (self.dim + 2)
530            ilevel1flat = ilevel1 + jnp.arange(x.shape[0]) * (self.dim + 2)
531
532            out = (
533                jnp.zeros((x.shape[0] * (self.dim + 2)), dtype=x.dtype)
534                .at[ilevelflat]
535                .set(w)
536                .at[ilevel1flat]
537                .set(1 - w)
538                .reshape(-1, self.dim + 2)[:, 1:-1]
539            )
540
541        elif basis == "exp_lr" or basis == "exp":
542            zeta = self.extra_params.get("zeta", default=2.0)
543            s = self.extra_params.get("s", default=0.5)
544            n = np.arange(self.dim)
545            # if self.trainable:
546            # zeta = jnp.abs(
547            #     self.param("zeta", lambda key: jnp.asarray(zeta, dtype=x.dtype))
548            # )
549            # s = jnp.abs(self.param("s", lambda key: jnp.asarray(s, dtype=x.dtype)))
550
551            a = zeta * s**n
552            xx = np.linspace(0, self.end, 10000)
553
554            if self.start > 0:
555                a1 = np.minimum(1.0 / a, 1.0)
556                switchstart = jnp.where(
557                    x[:, None] < self.end,
558                    1
559                    - (
560                        0.5
561                        + 0.5
562                        * jnp.cos(
563                            np.pi * (x[:, None] - self.start) / (self.end - self.start)
564                        )
565                    )
566                    ** a1[None, :],
567                    1,
568                ) * (x[:, None] > self.start)
569                switchstartxx = (
570                    1
571                    - (
572                        0.5
573                        + 0.5
574                        * np.cos(
575                            np.pi * (xx[:, None] - self.start) / (self.end - self.start)
576                        )
577                    )
578                    ** a1[None, :]
579                ) * (xx[:, None] > self.start)
580
581            else:
582                switchstart = 1.0
583                switchstartxx = 1.0
584
585            norm = 1.0 / np.trapz(
586                switchstartxx * np.exp(-a[None, :] * xx[:, None]), xx, axis=0
587            )
588            # norm = 1./np.max(switchstartxx*np.exp(-a[None, :] * xx[:, None]), axis=0)
589
590            if self.trainable:
591                a = jnp.abs(self.param("exponent", lambda key: jnp.asarray(a)))
592
593            out = switchstart * jnp.exp(-a[None, :] * x[:, None]) * norm[None, :]
594
595        elif basis == "neural_net" or basis == "nn":
596            neurons = self.extra_params.get(
597                "hidden_neurons", default=[2 * self.dim]
598            ) + [self.dim]
599            activation = self.extra_params.get("activation", default="swish")
600            use_bias = self.extra_params.get("use_bias", default=True)
601
602            out = FullyConnectedNet(
603                neurons, activation=activation, use_bias=use_bias, squeeze=False
604            )(x[:, None])
605        elif basis == "damped_coulomb":
606            l = np.arange(self.dim)[None, :]
607            a2 = self.extra_params.get("a", default=1.0)**2
608            x = x[:, None] - self.start
609            end = self.end - self.start
610            x21 = a2 + x**2
611            R21 = a2 + end**2
612            out = (
613                1.0 / x21 ** (0.5 * (l + 1))
614                - 1.0 / (R21 ** (0.5 * (l + 1)))
615                + (x - end) * ((l + 1) * end / (R21 ** (0.5 * (l + 3))))
616            ) * (x < end)
617
618        elif basis.startswith("spherical_bessel_"):
619            l = int(basis.split("_")[-1])
620            out = generate_spherical_jn_basis([l], self.dim, self.end)(x)
621        else:
622            raise NotImplementedError(f"Unknown radial basis {basis}.")
623        ############################
624
625        out = out.reshape((*shape, self.dim))
626
627        if self.graph_key is not None:
628            output_key = self.name if self.output_key is None else self.output_key
629            return {**inputs, output_key: out}
630        return out
631
632
633def positional_encoding_static(t, d: int, n: float = 10000.0):
634    if d % 2 == 0:
635        k = np.arange(d // 2)
636    else:
637        k = np.arange((d + 1) // 2)
638    wk = np.asarray(1.0 / (n ** (2 * k / d)))
639    wkt = wk[None, :] * t[:, None]
640    out = np.concatenate([np.cos(wkt), np.sin(wkt)], axis=-1)
641    if d % 2 == 1:
642        out = out[:, :-1]
643    return out
644
645
646@partial(jax.jit, static_argnums=(1, 2), inline=True)
647def positional_encoding(t, d: int, n: float = 10000.0):
648    if d % 2 == 0:
649        k = np.arange(d // 2)
650    else:
651        k = np.arange((d + 1) // 2)
652    wk = jnp.asarray(1.0 / (n ** (2 * k / d)))
653    wkt = wk[None, :] * t[:, None]
654    out = jnp.concatenate([jnp.cos(wkt), jnp.sin(wkt)], axis=-1)
655    if d % 2 == 1:
656        out = out[:, :-1]
657    return out
658
659
660def generate_spherical_jn_basis(dim:int, rc:float, ls:Union[int,Sequence[int]]=[0], print_code:bool=False, jit:bool=False):
661    from sympy import Symbol, jn, expand_func
662    from scipy.special import spherical_jn
663    from sympy import jn_zeros
664    import scipy.integrate as integrate
665
666    if isinstance(ls, int):
667        ls = list(range(ls + 1))
668    zl = [Symbol(f"xz[...,{l}]") for l in ls]
669    zn = np.array([jn_zeros(l, dim) for l in ls], dtype=float).T
670    znrc = zn / rc
671    norms = np.zeros((dim, len(ls)), dtype=float)
672    for l in ls:
673        for i in range(dim):
674            norms[i, l] = (
675                integrate.quad(lambda x: (spherical_jn(l, x) * x) ** 2, 0, zn[i, l])[0]
676                / znrc[i, l] ** 3
677            ) ** (-0.5)
678
679    fn_str = f"""def spherical_jn_basis_(x):
680    from jax.numpy import cos,sin
681    
682    znrc = jnp.array({znrc.tolist()},dtype=x.dtype)
683    norms = jnp.array({norms.tolist()},dtype=x.dtype)
684    xshape = x.shape
685    x = x.reshape(-1)
686    xz = x[:,None,None]*znrc[None,:,:]
687
688    jns = jnp.stack([
689  """
690    for l in ls:
691        fn_str += f"      {expand_func(jn(l, zl[l]))},\n"
692    fn_str += f"""    ],axis=-1)
693    return (norms[None,:,:]*jns).reshape(*xshape,{dim},{len(ls)})
694  """
695
696    if print_code:
697        print(fn_str)
698    exec(fn_str)
699    jn_basis = locals()["spherical_jn_basis_"]
700    if jit:
701        jn_basis = jax.jit(jn_basis)
702    return jn_basis