Discovering multiscale deep formulas in complex systems via neural-guided lambda calculus.

1. """
2. Main Deflex estimator implementing the three-stage training process.
3. """
5. import numpy as np
6. from sklearn.base import BaseEstimator, RegressorMixin
7. from sklearn.utils.validation import check_X_y, check_array
8. from typing import Optional, Dict, Any, Union
10. from ..deflexformer import DeflexformerEstimator
11. from ..lambda_regression import LambdaRegressionEstimator
14. class DeflexEstimator(BaseEstimator, RegressorMixin):
15. """
16. Deep Formula Discovery estimator combining neural networks and symbolic regression.
18. This estimator implements a three-stage training process:
19. 1. Pre-training: Train a single Deflexformer block on synthetic data
20. 2. Post-training: Stack pre-trained blocks and fine-tune on experimental data
21. 3. Symbolic regression: Extract symbolic formulas from the trained model
23. Parameters
24. ----------
25. n_blocks : int, default=5
26. Number of Deflexformer blocks to use in the model.
27. embedding_dim : int, default=64
28. Maximum embedding dimension for the Deflexformer.
29. n_synthetic_samples : int, default=10000
30. Number of synthetic samples to generate for pre-training.
31. pretrain_epochs : int, default=100
32. Number of epochs for pre-training stage.
33. finetune_epochs : int, default=50
34. Number of epochs for post-training/fine-tuning stage.
35. operator_set : list, optional
36. Set of operators for symbolic regression. If None, uses default set.
37. lambda_regression_params : dict, optional
38. Additional parameters for the LambdaRegression estimator.
39. random_state : int, optional
40. Random seed for reproducibility.
42. Attributes
43. ----------
44. deflexformer_ : DeflexformerEstimator
45. The trained Deflexformer model.
46. lambda_regression_ : LambdaRegressionEstimator
47. The symbolic regression model.
48. is_fitted_ : bool
49. Whether the estimator has been fitted.
50. formula_ : str
51. The discovered symbolic formula (available after fitting).
52. """
54. def __init__(
55. self,
56. n_blocks: int = 5,
57. embedding_dim: int = 64,
58. n_synthetic_samples: int = 10000,
59. pretrain_epochs: int = 100,
60. finetune_epochs: int = 50,
61. operator_set: Optional[list] = None,
62. lambda_regression_params: Optional[Dict[str, Any]] = None,
63. random_state: Optional[int] = None
64. ):
65. self.n_blocks = n_blocks
66. self.embedding_dim = embedding_dim
67. self.n_synthetic_samples = n_synthetic_samples
68. self.pretrain_epochs = pretrain_epochs
69. self.finetune_epochs = finetune_epochs
70. self.operator_set = operator_set
71. self.lambda_regression_params = lambda_regression_params or {}
72. self.random_state = random_state
74. def fit(self, X: np.ndarray, y: np.ndarray) -> 'DeflexEstimator':
75. """
76. Fit the Deflex model using the three-stage training process.
78. Parameters
79. ----------
80. X : array-like of shape (n_frames, n_elements, embedding_dim)
81. Training data as 3D tensor.
82. y : array-like of shape (n_frames, n_elements, output_dim)
83. Target values.
85. Returns
86. -------
87. self : DeflexEstimator
88. Fitted estimator.
89. """
90. # Validate input
91. X, y = check_X_y(X, y, multi_output=True, allow_nd=True)
93. if X.ndim != 3:
94. raise ValueError(f"X must be 3D tensor, got shape {X.shape}")
95. if y.ndim < 2:
96. raise ValueError(f"y must be at least 2D, got shape {y.shape}")
98. # Stage 1: Pre-training on synthetic data
99. self._pretrain_stage()
101. # Stage 2: Post-training on experimental data
102. self._posttrain_stage(X, y)
104. # Stage 3: Symbolic regression
105. self._symbolic_regression_stage(X, y)
107. self.is_fitted_ = True
108. return self
110. def _pretrain_stage(self) -> None:
111. """Stage 1: Pre-train single block on synthetic data."""
112. print("Stage 1: Pre-training on synthetic data...")
114. # Initialize LambdaRegression for synthetic data generation
115. self.lambda_regression_ = LambdaRegressionEstimator(
116. operator_set=self.operator_set,
117. random_state=self.random_state,
118. **self.lambda_regression_params
119. )
121. # Generate synthetic data
122. X_synthetic, y_synthetic = self.lambda_regression_.generate_synthetic_data(
123. n_samples=self.n_synthetic_samples
124. )
126. # Initialize single-block Deflexformer for pre-training
127. self.deflexformer_ = DeflexformerEstimator(
128. n_blocks=1, # Only one block for pre-training
129. embedding_dim=self.embedding_dim,
130. enable_clipping=True, # Strict clipping for synthetic data
131. random_state=self.random_state
132. )
134. # Pre-train the single block
135. self.deflexformer_.fit(
136. X_synthetic, y_synthetic,
137. epochs=self.pretrain_epochs,
138. stage="pretrain"
139. )
141. def _posttrain_stage(self, X: np.ndarray, y: np.ndarray) -> None:
142. """Stage 2: Stack blocks and fine-tune on experimental data."""
143. print("Stage 2: Post-training on experimental data...")
145. # Expand to full architecture with pre-trained weights
146. self.deflexformer_.expand_to_full_model(self.n_blocks)
148. # Fine-tune on experimental data
149. self.deflexformer_.fit(
150. X, y,
151. epochs=self.finetune_epochs,
152. stage="posttrain"
153. )
155. def _symbolic_regression_stage(self, X: np.ndarray, y: np.ndarray) -> None:
156. """Stage 3: Extract symbolic formulas."""
157. print("Stage 3: Symbolic regression...")
159. # Extract features from trained Deflexformer blocks
160. block_outputs = self.deflexformer_.extract_block_outputs(X)
162. # Generate additional synthetic data for symbolic regression
163. X_additional, y_additional = self.lambda_regression_.generate_synthetic_data(
164. n_samples=self.n_synthetic_samples // 2
165. )
167. # Combine experimental and synthetic data for symbolic regression
168. X_combined = np.concatenate([block_outputs, X_additional], axis=0)
169. y_combined = np.concatenate([y, y_additional], axis=0)
171. # Perform symbolic regression
172. self.lambda_regression_.fit(X_combined, y_combined)
173. self.formula_ = self.lambda_regression_.get_formula()
175. def predict(self, X: np.ndarray) -> np.ndarray:
176. """
177. Predict using the fitted model.
179. Parameters
180. ----------
181. X : array-like of shape (n_frames, n_elements, embedding_dim)
182. Input data.
184. Returns
185. -------
186. y_pred : ndarray of shape (n_frames, n_elements, output_dim)
187. Predicted values.
188. """
189. if not hasattr(self, 'is_fitted_'):
190. raise ValueError("This DeflexEstimator instance is not fitted yet.")
192. X = check_array(X, allow_nd=True)
194. # Use symbolic formula if available, otherwise use Deflexformer
195. if hasattr(self, 'formula_') and self.formula_:
196. return self.lambda_regression_.predict(X)
197. else:
198. return self.deflexformer_.predict(X)
200. def get_formula(self) -> str:
201. """
202. Get the discovered symbolic formula.
204. Returns
205. -------
206. formula : str
207. The symbolic formula in human-readable format.
208. """
209. if not hasattr(self, 'formula_'):
210. raise ValueError("Model must be fitted before getting formula.")
211. return self.formula_
213. def score(self, X: np.ndarray, y: np.ndarray) -> float:
214. """
215. Return the coefficient of determination R^2 of the prediction.
217. Parameters
218. ----------
219. X : array-like of shape (n_frames, n_elements, embedding_dim)
220. Test samples.
221. y : array-like of shape (n_frames, n_elements, output_dim)
222. True values.
224. Returns
225. -------
226. score : float
227. R^2 coefficient of determination.
228. """
229. from sklearn.metrics import r2_score
230. y_pred = self.predict(X)
231. return r2_score(y, y_pred, multioutput='uniform_average')

estimator.py at commit df664b1, under MIT · at the source