| import os, sys, math, struct |
| os.environ['JAX_PLATFORMS'] = 'cpu' |
| import jax |
| import jax.numpy as jnp |
| import flax.nnx as nn |
| import numpy as np |
| from typing import Tuple |
|
|
| print(f"=== CogniThread SSM Architecture === JAX {jax.__version__}") |
|
|
| |
| |
| |
| class WaveScalar(nn.Module): |
| """ |
| Encodes each scalar as a radionics scalar wave with 5 parameters: |
| amplitude, phase_shift, frequency, period (1/freq), fractal_dim. |
| v(t) = A * sin(2*pi*f*t + phi) * fractal_mod(t, D) |
| fractal_mod = 1 + sum_{n=1}^{N} (1/n^D) * sin(2*pi*n*f*t + n*phi) |
| """ |
| def __init__(self, shape: Tuple[int, ...], *, rngs: nn.Rngs, |
| init_amp_scale: float = 1.0, |
| freq_range: Tuple[float, float] = (0.1, 10.0), |
| fractal_range: Tuple[float, float] = (1.0, 2.0)): |
| super().__init__() |
| self.shape = shape |
| key = rngs.params() |
| fan_in = max(math.prod(shape), 1) |
| std = init_amp_scale / math.sqrt(fan_in) |
| self.amplitude = nn.Param(jax.random.normal(key, shape) * std) |
| key, sk = jax.random.split(key) |
| self.phase_shift = nn.Param(jax.random.uniform(sk, shape, minval=0, maxval=2*jnp.pi)) |
| key, sk = jax.random.split(key) |
| lf = jax.random.uniform(sk, shape, minval=jnp.log(freq_range[0]), maxval=jnp.log(freq_range[1])) |
| self.frequency = nn.Param(jnp.exp(lf)) |
| key, sk = jax.random.split(key) |
| self.fractal_dim = nn.Param(jax.random.uniform(sk, shape, minval=fractal_range[0], maxval=fractal_range[1])) |
|
|
| @property |
| def period(self): |
| return 1.0 / self.frequency.value |
|
|
| def fractal_mod(self, t, D, n_harm=8): |
| harms = jnp.arange(1, n_harm+1, dtype=jnp.float32) |
| w = 1.0 / jnp.power(harms, D[..., None]) |
| f, phi = self.frequency.value, self.phase_shift.value |
| sigs = jnp.sin(2*jnp.pi*harms*f[..., None]*t + harms*phi[..., None]) |
| mod = 1.0 + jnp.sum(w * sigs, axis=-1) |
| return mod / jnp.sqrt(1.0 + jnp.sum(w**2, axis=-1)) |
|
|
| def evaluate(self, t=None): |
| if t is None: |
| t = jnp.array(0.0) |
| A = self.amplitude.value |
| f = self.frequency.value |
| phi = self.phase_shift.value |
| D = self.fractal_dim.value |
| base = A * jnp.sin(2*jnp.pi*f*t + phi) |
| return base * self.fractal_mod(t, D) |
|
|
| def decode(self): |
| return self.evaluate(t=jnp.array(0.0)) |
|
|
| def rms_value(self): |
| A, D = self.amplitude.value, self.fractal_dim.value |
| harms = jnp.arange(1, 9, dtype=jnp.float32) |
| w = 1.0 / jnp.power(harms, D[..., None]) |
| rf = jnp.sqrt(1.0 + jnp.sum(w**2, axis=-1)/2.0) |
| nm = jnp.sqrt(1.0 + jnp.sum(w**2, axis=-1)) |
| return A / jnp.sqrt(2.0) * rf / nm |
|
|
|
|
| class WaveScalarLinear(nn.Module): |
| def __init__(self, in_f: int, out_f: int, *, rngs: nn.Rngs): |
| super().__init__() |
| self.weight = WaveScalar((out_f, in_f), rngs=rngs) |
| self.bias = WaveScalar((out_f,), rngs=rngs, init_amp_scale=0.01) |
|
|
| def __call__(self, x): |
| return x @ self.weight.decode().T + self.bias.decode() |
|
|
| print("[1/6] WaveScalar defined") |
|
|
