""" starG_helpers.py Utility functions for ★_G comparison experiments LH & SU & Claude 2026 """ import numpy as np from scipy.spatial.distance import pdist from typing import List, Tuple, Union class StarGHelpers: """ Utility functions for ★_G comparison experiments. All methods are static and can be called without instantiation. """ @staticmethod def compute_r2(y_pred: np.ndarray, y_true: np.ndarray) -> float: """ Compute R² (coefficient of determination) score. Parameters ---------- y_pred : array-like Predicted values y_true : array-like True values Returns ------- r2 : float R² score (1.0 is perfect prediction) """ y_pred = np.asarray(y_pred).flatten() y_true = np.asarray(y_true).flatten() ss_res = np.sum((y_pred - y_true) ** 2) ss_tot = np.sum((y_true - np.mean(y_true)) ** 2) + 1e-10 r2 = 1 - ss_res / ss_tot return r2 @staticmethod def count_mlp_params(layers: List[int]) -> int: """ Count total parameters in a Multi-Layer Perceptron. Parameters ---------- layers : list of int Number of neurons in each layer (including input and output) Returns ------- n_params : int Total number of trainable parameters (weights + biases) Example ------- >>> StarGHelpers.count_mlp_params([784, 256, 128, 10]) 235146 # (784*256 + 256) + (256*128 + 128) + (128*10 + 10) """ n_params = 0 for l in range(len(layers) - 1): # Weights: layers[l] * layers[l+1] # Biases: layers[l+1] n_params += layers[l] * layers[l + 1] + layers[l + 1] return n_params @staticmethod def compute_invariant_features(X: np.ndarray) -> np.ndarray: """ Compute rotation-invariant features from 3D tensor. Computes statistics (mean, std, min, max) across the group dimension (axis 2) to create rotation-invariant representations. Parameters ---------- X : ndarray of shape (n_samples, n_features, n_rotations) Input tensor with group structure Returns ------- X_inv : ndarray of shape (n_samples, 4 * n_features) Invariant features: [mean, std, min, max] concatenated """ X_mean = np.mean(X, axis=2) X_std = np.std(X, axis=2, ddof=0) # ddof=0 matches MATLAB's default X_min = np.min(X, axis=2) X_max = np.max(X, axis=2) X_inv = np.concatenate([X_mean, X_std, X_min, X_max], axis=1) return X_inv @staticmethod def generate_molecular_data( n_samples: int, n_feat: int, n_rot: int ) -> Tuple[np.ndarray, np.ndarray]: """ Generate synthetic molecular data with rotational symmetry. Creates molecules with random atom positions and atomic numbers, then generates features under different rotations around the z-axis. Parameters ---------- n_samples : int Number of molecules to generate n_feat : int Number of features per rotation n_rot : int Number of rotations (group order) Returns ------- X : ndarray of shape (n_samples, n_feat, n_rot) Feature tensor with rotational structure Y : ndarray of shape (n_samples, 1) Target values (rotation-invariant molecular property) Notes ----- The target Y is computed from rotation-invariant quantities: - Mean pairwise distance between atoms - Standard deviation of pairwise distances - Sum of squared atomic numbers Features include: - Weighted atomic positions (x, y, z) - Weighted radial distances - Angular harmonics with Gaussian decay """ X = np.zeros((n_samples, n_feat, n_rot)) Y = np.zeros((n_samples, 1)) for i in range(n_samples): # Generate random molecule n_atoms = np.random.randint(4, 11) # [4, 10] inclusive pos = np.random.randn(n_atoms, 3) * 2 pos = pos - np.mean(pos, axis=0) # Center at origin # Random atomic numbers, biased toward carbon (6) Z = np.random.randint(1, 10, n_atoms) # [1, 9] inclusive Z[(Z > 1) & (Z < 6)] = 6 # Replace 2,3,4,5 with 6 (carbon) # Compute rotation-invariant target dists = pdist(pos) if len(dists) == 0: dists = np.array([1.0]) Y[i] = np.mean(dists) + 0.3 * np.std(dists) + np.sum(Z ** 2) / 500 # Generate features for each rotation for g in range(n_rot): theta = 2 * np.pi * g / n_rot # Rotation matrix around z-axis R = np.array([ [np.cos(theta), -np.sin(theta), 0], [np.sin(theta), np.cos(theta), 0], [0, 0, 1] ]) # Rotate positions pos_rot = (R @ pos.T).T # Compute features feat = np.zeros(n_feat) for a in range(n_atoms): x, y, z = pos_rot[a] r = np.linalg.norm(pos_rot[a]) # Basic weighted features feat[0] += Z[a] * x feat[1] += Z[a] * y feat[2] += Z[a] * z feat[3] += Z[a] * r feat[4] += Z[a] * r ** 2 # Angular harmonics with Gaussian decay for f in range(5, n_feat): angular_freq = f - 4 # 1, 2, 3, ... angle = np.arctan2(y, x) feat[f] += Z[a] * np.cos(angular_freq * angle) * np.exp(-r ** 2 / 8) X[i, :, g] = feat return X, Y @staticmethod def roundp(x: Union[float, np.ndarray], p: int) -> Union[float, np.ndarray]: """ Round to specified number of decimal places. Parameters ---------- x : float or ndarray Value(s) to round p : int Number of decimal places Returns ------- rounded : float or ndarray Rounded value(s) Example ------- >>> StarGHelpers.roundp(3.14159, 2) 3.14 >>> StarGHelpers.roundp(np.array([1.234, 5.678]), 1) array([1.2, 5.7]) """ scale = 10 ** p return np.round(x * scale) / scale # ============================================================================= # Convenience functions (for direct import) # ============================================================================= def compute_r2(y_pred, y_true): """Alias for StarGHelpers.compute_r2""" return StarGHelpers.compute_r2(y_pred, y_true) def count_mlp_params(layers): """Alias for StarGHelpers.count_mlp_params""" return StarGHelpers.count_mlp_params(layers) def compute_invariant_features(X): """Alias for StarGHelpers.compute_invariant_features""" return StarGHelpers.compute_invariant_features(X) def generate_molecular_data(n_samples, n_feat, n_rot): """Alias for StarGHelpers.generate_molecular_data""" return StarGHelpers.generate_molecular_data(n_samples, n_feat, n_rot) def roundp(x, p): """Alias for StarGHelpers.roundp""" return StarGHelpers.roundp(x, p) # ============================================================================= # Example Usage and Tests # ============================================================================= if __name__ == "__main__": print("=" * 60) print("StarGHelpers Test Suite") print("=" * 60) # Test compute_r2 print("\n1. Testing compute_r2:") y_true = np.array([1.0, 2.0, 3.0, 4.0, 5.0]) y_pred_perfect = y_true.copy() y_pred_good = y_true + np.random.randn(5) * 0.1 y_pred_bad = np.ones(5) * np.mean(y_true) print(f" Perfect prediction R²: {compute_r2(y_pred_perfect, y_true):.4f}") print(f" Good prediction R²: {compute_r2(y_pred_good, y_true):.4f}") print(f" Mean prediction R²: {compute_r2(y_pred_bad, y_true):.4f}") # Test count_mlp_params print("\n2. Testing count_mlp_params:") architectures = [ [10, 20, 1], [784, 256, 128, 10], [100, 64, 32, 16, 1] ] for arch in architectures: n_params = count_mlp_params(arch) print(f" Architecture {arch}: {n_params:,} parameters") # Test compute_invariant_features print("\n3. Testing compute_invariant_features:") X_test = np.random.randn(5, 8, 4) # 5 samples, 8 features, 4 rotations X_inv = compute_invariant_features(X_test) print(f" Input shape: {X_test.shape}") print(f" Output shape: {X_inv.shape}") print(f" Expected: ({X_test.shape[0]}, {4 * X_test.shape[1]})") # Test generate_molecular_data print("\n4. Testing generate_molecular_data:") n_samples, n_feat, n_rot = 100, 16, 8 X, Y = generate_molecular_data(n_samples, n_feat, n_rot) print(f" Generated X shape: {X.shape}") print(f" Generated Y shape: {Y.shape}") print(f" Y statistics: min={Y.min():.3f}, max={Y.max():.3f}, mean={Y.mean():.3f}") # Verify rotational equivariance of features print("\n Checking rotational structure:") # For cyclic features, check that feature magnitude varies smoothly sample_idx = 0 feat_magnitudes = np.linalg.norm(X[sample_idx], axis=0) print(f" Feature magnitudes across rotations: {feat_magnitudes}") # Test roundp print("\n5. Testing roundp:") test_values = [3.14159265, 2.71828182, 1.41421356] for val in test_values: print(f" {val} -> p=2: {roundp(val, 2)}, p=4: {roundp(val, 4)}") # Test with array arr = np.array([1.234567, 8.901234]) print(f" Array {arr} -> p=3: {roundp(arr, 3)}") print("\n" + "=" * 60) print("All tests completed successfully!") print("=" * 60)