File size: 3,318 Bytes
350f35d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 | 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__}")
###############################################################################
# 1. WAVESCALAR — Thermal Analog Wave Representation
###############################################################################
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")
|