import numpy as np from scipy.io.wavfile import read # Load training and test data X_train = read('../X_train.wav')[1].reshape(-1, 784) * 2 - 1 y_train = (read('../y_train.wav')[1] * 9).astype(int) X_test = read('../X_test.wav')[1].reshape(-1, 784) * 2 - 1 y_test = (read('../y_test.wav')[1] * 9).astype(int) X20 = X_test[:1000] yt20 = y_test[:1000] class CCODE_MLP_Fast: """ Fast CCT-ODE Classifier using vectorised periodic layers and SPSA training. """ def __init__(self, input_size, hidden_size, output_size, n_functions=4): self.input_size = input_size self.hidden_size = hidden_size self.output_size = output_size self.n_functions = n_functions # Trainable parameters (minimal) self.phases = np.random.rand(n_functions) * 2 * np.pi self.frequencies = np.random.rand(n_functions) * 2.0 self.amplitudes = np.random.rand(n_functions) * 0.1 # Output bias (small periodic offset) self.output_bias = np.random.rand(output_size) * 0.01 # Precompute frequency multipliers for each input dimension # freq_ij = base_freq * (1 + 0.01*j) for j in 0..input_size-1 j_idx = np.arange(input_size, dtype=np.float32) self.freq_multiplier = 1.0 + 0.01 * j_idx # shape (input_size,) # Precompute feature indices for output layer (modular reuse) self.output_func_idx = np.arange(output_size) % n_functions # ODE state variables self.t = 0 self.H = 1.0 self.collapse_threshold = 0.1 self.state_history = [] # ---------- Periodic functions (vectorised) ---------- @staticmethod def sine_wave(angle, amp): """amp * sin(angle) on whole arrays.""" return amp * np.sin(angle) @staticmethod def triangle_wave(x, freq, phase, amp): """Vectorised triangle wave (for output layer).""" y = (x * freq + phase / (2 * np.pi)) % 1.0 tri = 2 * np.abs(2 * y - 1) - 1 return amp * tri # ---------- Fast forward pass ---------- def periodic_forward(self, x): """ Vectorised forward pass. x : (batch, input_size) returns probs : (batch, output_size) """ batch_size = x.shape[0] # Normalise input per sample x_norm = x / (np.abs(x).max(axis=1, keepdims=True) + 1e-9) # ----- Hidden layer: sine waves (fully vectorised) ----- # Build matrices: (hidden_size, input_size) for base freq/phase/amp # Each hidden neuron i uses func_idx = i % n_functions hid_indices = np.arange(self.hidden_size) func_idx_hid = hid_indices % self.n_functions # Retrieve base parameters as (hidden_size,) arrays base_freq = self.frequencies[func_idx_hid] # (H,) base_phase = self.phases[func_idx_hid] # (H,) base_amp = self.amplitudes[func_idx_hid] # (H,) # Expand to (H, D) using precomputed freq_multiplier # freq_ij = base_freq_i * (1 + 0.01*j) freq_mat = np.outer(base_freq, self.freq_multiplier) # (H, D) # Phase matrix: base_phase_i + x_j * π -> depends on input x # We'll compute angle = 2π * freq_ij * x_j + phase_ij # phase_ij = base_phase_i + π * x_j # Hence angle = 2π * freq_ij * x_j + base_phase_i + π * x_j # First term: 2π * (freq_mat * x) -> (B, H, D) ? careful with shapes # Better: compute for each sample using broadcasting: # x_norm shape (B, D) , freq_mat shape (H, D) # term1 = 2π * (x_norm * freq_mat) -> (B, H, D) after broadcasting # term2 = x_norm * π -> (B, D) # term3 = base_phase -> (H,) # Using Einstein summation for clarity: # angle[b, h, d] = 2π * (x_norm[b, d] * freq_mat[h, d]) + base_phase[h] + π * x_norm[b, d] # Then sum over d. # Compute contributions: for each d, for each h, amplitude * sin(angle) # We can do it step by step with broadcasting: # shape (B, 1, D) * (1, H, D) -> (B, H, D) for product x_exp = x_norm[:, None, :] # (B, 1, D) freq_exp = freq_mat[None, :, :] # (1, H, D) angle_1 = 2 * np.pi * x_exp * freq_exp # (B, H, D) # Add π * x angle_2 = np.pi * x_exp # (B, 1, D) angle_3 = base_phase[None, :, None] # (1, H, 1) angle = angle_1 + angle_2 + angle_3 # (B, H, D) # Sine and amplitude sin_vals = np.sin(angle) # (B, H, D) amp_exp = base_amp[None, :, None] # (1, H, 1) contributions = amp_exp * sin_vals # (B, H, D) # Sum over input dimension -> hidden state (B, H) hidden = np.sum(contributions, axis=2) # Bias via sawtooth (simplified: using sine) bias = self.sine_wave(self.t * 0.5 + self.output_bias[hid_indices % self.output_size], 0.01) hidden += bias[None, :] # Activation: tanh (bounded oscillator) hidden = np.tanh(hidden) # ----- Output layer: triangle waves (loop over output classes) ----- output = np.zeros((batch_size, self.output_size)) for i in range(self.output_size): func_idx = self.output_func_idx[i] freq_i = self.frequencies[func_idx] phase_i = self.phases[func_idx] amp_i = self.amplitudes[func_idx] * 0.5 # Sum over hidden units for h in range(self.hidden_size): output[:, i] += self.triangle_wave(hidden[:, h], freq_i, phase_i + h * 0.05, amp_i) # Periodic centering + softmax output_max = np.max(output, axis=1, keepdims=True) output = output - output_max exp_out = np.exp(output) probs = exp_out / np.sum(exp_out, axis=1, keepdims=True) # Entropy (only if needed – here we keep for compatibility) self.H = -np.mean(np.sum(probs * np.log(probs + 1e-9), axis=1)) return probs # ---------- Periodicity detection (unchanged, but called less often) ---------- def detect_periodicity(self, state): self.state_history.append(state.copy()) if len(self.state_history) > 20: self.state_history.pop(0) if len(self.state_history) >= 10: for k in range(1, 5): if len(self.state_history) >= 2*k: if np.allclose(self.state_history[-k], self.state_history[-2*k], atol=0.1): return k return 0 # ---------- SPSA training (fast gradient estimation) ---------- def _pack_params(self): """Concatenate all trainable parameters into a flat array.""" return np.concatenate([self.phases, self.frequencies, self.amplitudes]) def _unpack_params(self, theta): """Restore parameters from flat array.""" n = self.n_functions self.phases = theta[:n] self.frequencies = theta[n:2*n] self.amplitudes = theta[2*n:3*n] def _loss(self, x, y_true, theta): """Compute cross-entropy loss for given parameters.""" self._unpack_params(theta) y_pred = self.periodic_forward(x) loss = -np.mean(np.sum(y_true * np.log(y_pred + 1e-9), axis=1)) return loss def update(self, x, y_true, learning_rate=0.01, spsa_c=0.01): """ One SPSA update step. Uses only 2 extra forward passes. """ # Current parameters theta = self._pack_params() n_params = len(theta) # Generate random perturbation vector (Rademacher distribution) delta = np.random.choice([-1, 1], size=n_params) # Compute loss at theta + c*delta and theta - c*delta loss_plus = self._loss(x, y_true, theta + spsa_c * delta) loss_minus = self._loss(x, y_true, theta - spsa_c * delta) # Estimate gradient grad_est = (loss_plus - loss_minus) / (2 * spsa_c) * delta # Update parameters theta_new = theta - learning_rate * grad_est self._unpack_params(theta_new) # Periodic constraints self.phases = self.phases % (2 * np.pi) self.amplitudes = np.abs(self.amplitudes) + 0.001 # Time step self.t += 1 # Return current prediction (for logging) return self.periodic_forward(x) def predict(self, x): probs = self.periodic_forward(x) return np.argmax(probs, axis=1) def score(self, x, y): acc = np.mean(self.predict(x) == y) print(f"Entropy: {self.H:.4f} | Accuracy: {acc:.4f}") return acc # ---------------------- DEMO ---------------------- print("=" * 60) print("FAST CCT-ODE Classifier (vectorised + SPSA)") print("=" * 60) # Create model cct_fast = CCODE_MLP_Fast(input_size=784, hidden_size=100, output_size=10, n_functions=4) print(f"\nWeight reduction:") print(f" Standard MLP: ~79,400 weights") print(f" Fast CCT-ODE: {cct_fast.n_functions * 3} parameters") print("\nTraining (fast mode)...") for epoch in range(50): # fewer epochs because each step is now fast idx = np.random.randint(0, 60000, 128) # larger batch for vectorisation X_batch = X_train[idx] y_batch = y_train[idx] y_onehot = np.eye(10)[y_batch] cct_fast.update(X_batch, y_onehot, learning_rate=0.05) if epoch % 1 == 0: acc = cct_fast.score(X_batch, y_batch) print(f" --- Epoch {epoch} done ---") print("\nTraining complete. Fast CCT-ODE ready.")