diff --git a/.github/workflows/pull-requests.yml b/.github/workflows/pull-requests.yml index a140156..7f6485c 100644 --- a/.github/workflows/pull-requests.yml +++ b/.github/workflows/pull-requests.yml @@ -31,6 +31,9 @@ jobs: - name: Lint with ruff run: >- ruff check tce/ + - name: Complexity check with complexipy + run: >- + complexipy tce/ --failed --suggest-refactors - name: Type check with mypy run: >- mypy tce/ diff --git a/.github/workflows/workflow.yml b/.github/workflows/workflow.yml index ff0745f..3b9b2e3 100644 --- a/.github/workflows/workflow.yml +++ b/.github/workflows/workflow.yml @@ -32,6 +32,9 @@ jobs: - name: Lint with ruff run: >- ruff check tce/ + - name: Complexity check with complexipy + run: >- + complexipy tce/ --failed --suggest-refactors - name: Type check with mypy run: >- mypy tce/ diff --git a/pyproject.toml b/pyproject.toml index da3f2ae..3b4841e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -49,6 +49,7 @@ dev = [ "pytest-cov~=7.0.0", "ruff~=0.9.4", "mypy~=1.13.0", + "complexipy~=7.0.1", "pdoc~=15.0.1", "scipy-stubs~=1.15.3.0", "scikit-learn~=1.7.1", diff --git a/tce/__init__.py b/tce/__init__.py index 29da033..fc1bd09 100644 --- a/tce/__init__.py +++ b/tce/__init__.py @@ -463,7 +463,7 @@ def callback(step_: int, num_steps_: int): """ -__version__ = "1.0.2" +__version__ = "1.0.3" __authors__ = ["Jacob Jeffries"] __url__ = "https://github.com/MUEXLY/tce-lib" diff --git a/tce/calculator.py b/tce/calculator.py index f50558b..589fb1f 100644 --- a/tce/calculator.py +++ b/tce/calculator.py @@ -216,6 +216,51 @@ class TCECalculator(Calculator): $m$ is the maximum body-order, and $\mathcal{A}$ denotes the set of elements included in the expansion. """ + def _register_many_body_feature( + self, + feature: tuple[int, ...], + seen_features: set[tuple[int, ...]] + ) -> None: + """Validate a many-body feature and store it in the correct body-order bucket.""" + + discriminant = 1 + 8 * len(feature) + integer_sqrt = isqrt(discriminant) + + if integer_sqrt * integer_sqrt != discriminant or (1 + integer_sqrt) % 2 != 0: + raise ValueError( + f"feature {feature} is invalid. Every feature label should have length (m choose 2) " + "where m is an integer" + ) + + canonical_feature = tuple(sorted(feature)) + if canonical_feature in seen_features: + warnings.warn(f"feature {feature} is a duplicate feature by symmetry") + + seen_features.add(canonical_feature) + self.feature_groups[(1 + integer_sqrt) // 2].append(feature) + + def _build_einsum_str(self, body_order: int) -> str: + """Build the einsum contraction string for a given body order.""" + + latin_indices = LATIN_ALPHABET[:body_order] + greek_indices = GREEK_ALPHABET[:body_order] + mixed_indices = ','.join( + f'{latin}{greek}' for latin, greek in zip(latin_indices, greek_indices) + ) + input_str = f"L{latin_indices},{mixed_indices}" + output_str = f"L{greek_indices}" + return f"{input_str}->{output_str}" + + def _compute_feature_vector_size(self) -> int: + """Pre-compute the size of the flattened feature vector.""" + + num_species = len(self.species) + feature_vector_size = len(self.neighbor_cutoffs) * (num_species ** 2) + for body_order in self.feature_groups.keys(): + if body_order >= 3: + feature_vector_size += len(self.feature_groups[body_order]) * (num_species ** body_order) + return feature_vector_size + def __post_init__(self): Calculator.__init__(self) @@ -232,52 +277,12 @@ def __post_init__(self): self.einsum_strs = {2: "Lij,iα,jβ->Lαβ"} seen_features: set[tuple[int, ...]] = set() for feature in self.many_body_features: + self._register_many_body_feature(feature, seen_features) - # each feature motif label is size (m choose 2) - # where m is the number of atoms in the motif - # we want to store m -> [list of features of size m] - # len(feature) here is the size of the label - # n = (m choose 2) implies that m^2 - m - 2n = 0 - # so m = (1 + sqrt(1 + 8n)) / 2 - - discriminant = 1 + 8 * len(feature) - integer_sqrt = isqrt(discriminant) - - if integer_sqrt * integer_sqrt != discriminant or (1 + integer_sqrt) % 2 != 0: - raise ValueError( - f"feature {feature} is invalid. Every feature label should have length (m choose 2) " - "where m is an integer" - ) - - if tuple(sorted(feature)) in seen_features: - warnings.warn(f"feature {feature} is a duplicate feature by symmetry") - - seen_features.add(tuple(sorted(feature))) - self.feature_groups[(1 + integer_sqrt) // 2].append(feature) - - # Pre-compute einsum strings for each body order - for body_order in self.feature_groups.keys(): - latin_indices = LATIN_ALPHABET[:body_order] - greek_indices = GREEK_ALPHABET[:body_order] - - # build mixed indices i\alpha,j\beta,k\gamma,etc - mixed_indices = ','.join( - f'{latin}{greek}' for latin, greek in zip(latin_indices, greek_indices) - ) - input_str = f"L{latin_indices},{mixed_indices}" - output_str = f"L{greek_indices}" - self.einsum_strs[body_order] = f"{input_str}->{output_str}" - - # Pre-compute feature vector size - num_species = len(self.species) - self.feature_vector_size = len(self.neighbor_cutoffs) * (num_species ** 2) - - # For body_order >= 3, the number of features is len(feature_groups[body_order]) for body_order in self.feature_groups.keys(): - if body_order >= 3: - num_features = len(self.feature_groups[body_order]) - self.feature_vector_size += num_features * (num_species ** body_order) + self.einsum_strs[body_order] = self._build_einsum_str(body_order) + self.feature_vector_size = self._compute_feature_vector_size() self.implemented_properties = list(self.models.keys()) @@ -307,13 +312,14 @@ def get_feature_label_order(self) -> list[tuple[tuple[int, ...], tuple[str, ...] for body_order, features in self.feature_groups.items(): if body_order < 3: continue - for feature in features: + + species_iterable = product(range(num_species), repeat=body_order) + for feature, species_indices in product(features, species_iterable): topological_label = tuple(sorted(feature)) - for species_indices in product(range(num_species), repeat=body_order): - species_multiset = tuple( - (self.species[idx] for idx in species_indices) - ) - labels.append((topological_label, species_multiset)) + species_multiset = tuple( + (self.species[idx] for idx in species_indices) + ) + labels.append((topological_label, species_multiset)) return labels @@ -341,52 +347,28 @@ def get_topological_tensors(self, atoms: Atoms) -> dict[int, sparse.COO]: topology_key = hash_topology(atoms) topological_tensors = self.topological_tensors.get(topology_key) - if topological_tensors is None: - - # these are boolean, so we can sum corresponding to logical or - adjacency_tensors = get_adjacency_tensors( - atoms=atoms, - cutoffs=self.neighbor_cutoffs, - tolerance=self.neighbor_tolerance - ) - topological_tensors = {2: adjacency_tensors} + if topological_tensors is not None: + return topological_tensors - for body_order, features in self.feature_groups.items(): - - final_result_str = LATIN_ALPHABET[:body_order] - input_str = ','.join( - f"{i1}{i2}" for i1, i2 in combinations(final_result_str, r=2) - ) - einsum_str = f"{input_str}->{final_result_str}" - - n_body_tensors = [] - for label in features: - n_body_tensor = sum( - contract( - einsum_str, - *(adjacency_tensors[letter] for letter in permuted_label) - ) for permuted_label in set(permutations(label)) - ) + # these are boolean, so we can sum corresponding to logical or + adjacency_tensors = get_adjacency_tensors( + atoms=atoms, + cutoffs=self.neighbor_cutoffs, + tolerance=self.neighbor_tolerance + ) + topological_tensors = {2: adjacency_tensors} - if not n_body_tensor.nnz: - warnings.warn(f"feature {label} is identically 0") - - n_body_tensors.append(n_body_tensor) - - stacked_tensors = sparse.stack(n_body_tensors) - difference = symmetrize(stacked_tensors, axes=tuple(range(1, 1 + body_order))) - stacked_tensors - if difference.nnz and not np.allclose(difference.data, 0): - raise ValueError( - f"Topological tensors for body order {body_order} are not symmetric in indices 1..{body_order}" - ) - topological_tensors[body_order] = stacked_tensors - - topological_tensors[2] = sparse.COO( - coords=topological_tensors[2].coords, - data=topological_tensors[2].data.astype(np.int64), - shape=topological_tensors[2].shape + for body_order, features in self.feature_groups.items(): + topological_tensors[body_order] = self._compute_topological_tensors_for_body_order( + body_order, features, adjacency_tensors ) - self.topological_tensors[topology_key] = topological_tensors + + topological_tensors[2] = sparse.COO( + coords=topological_tensors[2].coords, + data=topological_tensors[2].data.astype(np.int64), + shape=topological_tensors[2].shape + ) + self.topological_tensors[topology_key] = topological_tensors return topological_tensors @@ -753,6 +735,58 @@ def get_feature_vector_difference(self, initial: Atoms, final: Atoms) -> NDArray raise NotImplementedError + def _compute_topological_tensors_for_body_order( + self, + body_order: int, + features: list[tuple[int, ...]], + adjacency_tensors: dict[int, sparse.COO] + ) -> sparse.COO: + """Compute the symmetry-checked topological tensors for a single body order.""" + + final_result_str = LATIN_ALPHABET[:body_order] + input_str = ','.join( + f"{i1}{i2}" for i1, i2 in combinations(final_result_str, r=2) + ) + einsum_str = f"{input_str}->{final_result_str}" + + n_body_tensors = [] + for label in features: + n_body_tensor = sum( + contract( + einsum_str, + *(adjacency_tensors[letter] for letter in permuted_label) + ) for permuted_label in set(permutations(label)) + ) + + if not n_body_tensor.nnz: + warnings.warn(f"feature {label} is identically 0") + + n_body_tensors.append(n_body_tensor) + + stacked_tensors = sparse.stack(n_body_tensors) + difference = symmetrize(stacked_tensors, axes=tuple(range(1, 1 + body_order))) - stacked_tensors + if difference.nnz and not np.allclose(difference.data, 0): + raise ValueError( + f"Topological tensors for body order {body_order} are not symmetric in indices 1..{body_order}" + ) + return stacked_tensors + + def _calculate_property(self, atoms: Atoms, name: str) -> None: + """Compute and store one model property for a configuration.""" + + if name not in self.implemented_properties: + raise PropertyNotImplementedError(f"property {name} not included in models") + + feature_vec = self.get_feature_vector(atoms) + if self.intensive[name]: + feature_vec /= len(atoms) + + prop = self.models[name].predict(feature_vec.reshape(1, -1)) + if isinstance(prop, np.ndarray): + prop = prop.squeeze() + + self.results[name] = prop + def calculate( self, atoms: Optional[Atoms] = None, @@ -766,19 +800,7 @@ def calculate( raise ValueError for name in properties: - - if name not in self.implemented_properties: - raise PropertyNotImplementedError(f"property {name} not included in models") - - feature_vec = self.get_feature_vector(atoms) - if self.intensive[name]: - feature_vec /= len(atoms) - - prop = self.models[name].predict(feature_vec.reshape(1, -1)) - if isinstance(prop, np.ndarray): - prop = prop.squeeze() - - self.results[name] = prop + self._calculate_property(atoms, name) def train(self, configurations: list[Atoms]): diff --git a/tce/monte_carlo.py b/tce/monte_carlo.py index de8541b..4e6293e 100644 --- a/tce/monte_carlo.py +++ b/tce/monte_carlo.py @@ -46,6 +46,53 @@ def score(self, X, y): raise NotImplementedError +def _transform_sklearn_pipeline(model: Model) -> SurrogateModel: + + from_sklearn = model.__class__.__module__.startswith('sklearn') + if not from_sklearn: + raise ValueError("sklearn pipeline transformation is only valid for a model from sklearn") + + from sklearn.pipeline import Pipeline + + if not isinstance(model, Pipeline): + raise TypeError("model must be an instance of sklearn.pipeline.Pipeline") + def _infer_input_dim(pipeline: Pipeline): + # scan from the *front* + for _, step in pipeline.steps: + if hasattr(step, "n_features_in_"): + return step.n_features_in_ + + # fallback: try final estimator ONLY if nothing else exists + final = model.steps[-1][1] + if hasattr(final, "n_features_in_"): + return final.n_features_in_ + + raise ValueError("Could not infer input dimension.") + + d = _infer_input_dim(model) + + # find effective by plugging in basis vectors + + def _eval_pipeline(pipeline: Pipeline, X: NDArray) -> float: + y = pipeline.predict(X) + return float(np.asarray(y).reshape(-1)[0]) + + x0 = np.zeros((1, d)) + f0 = _eval_pipeline(model, x0) + + beta = np.zeros(d) + + # probe standard basis + for i in range(d): + xi = np.zeros((1, d)) + xi[0, i] = 1.0 + + fi = _eval_pipeline(model, xi) + beta[i] = fi - f0 + + return SurrogateModel(beta) + + def transform_model(model: Model) -> Model: r""" @@ -98,46 +145,55 @@ def transform_model(model: Model) -> Model: if isinstance(model, Pipeline): - # most complicated case, need to calculate an effective β + return _transform_sklearn_pipeline(model) - # find the final dimension d - def _infer_input_dim(pipeline): - # scan from the *front* - for _, step in pipeline.steps: - if hasattr(step, "n_features_in_"): - return step.n_features_in_ + raise NotImplementedError - # fallback: try final estimator ONLY if nothing else exists - final = pipeline.steps[-1][1] - if hasattr(final, "n_features_in_"): - return final.n_features_in_ - raise ValueError("Could not infer input dimension.") +def _default_mc_callback(step_: int, num_steps_: int): + LOGGER.info(f"MC step {step_:.0f}/{num_steps_:.0f}") - d = _infer_input_dim(model) - # find effective by plugging in basis vectors +def _default_mc_step(generator: np.random.Generator) -> Callable[[Atoms], Atoms]: + def _step(atoms: Atoms) -> Atoms: + new_atoms = atoms.copy() + i, j = generator.integers(len(atoms), size=2) + new_atoms[i].symbol, new_atoms[j].symbol = new_atoms[j].symbol, new_atoms[i].symbol + return new_atoms + return _step - def _eval_pipeline(pipeline, X): - y = pipeline.predict(X) - return float(np.asarray(y).reshape(-1)[0]) - x0 = np.zeros((1, d)) - f0 = _eval_pipeline(model, x0) +def _resolve_beta_values(beta: float | Sequence[float] | NDArray[np.floating], num_steps: int): + if isinstance(beta, (Sequence, np.ndarray)): + assert len(beta) == num_steps, "if beta is a sequence, it must be the same length as num_steps" + return np.asarray(beta, dtype=float) + if isinstance(beta, (float, int, np.floating)): + return np.full(num_steps, float(beta)) + raise TypeError("beta must be either a float or a sequence of floats") - beta = np.zeros(d) - # probe standard basis - for i in range(d): - xi = np.zeros((1, d)) - xi[0, i] = 1.0 +def _as_scalar(value): + if isinstance(value, np.ndarray): + return value.item() + return value - fi = _eval_pipeline(model, xi) - beta[i] = fi - f0 - return SurrogateModel(beta) +def _prepare_energy_model(tce_calculator: TCECalculator, initial_configuration: Atoms): + zero_feature = np.zeros(tce_calculator.feature_vector_size).reshape(1, -1) + predicted = _as_scalar(tce_calculator.models["energy"].predict(zero_feature)) - raise NotImplementedError + if predicted != 0.0: + warnings.warn( + "Input model has an intercept, which will mess with energy difference calculations. " + "The monte carlo run will automatically zero-out this intercept, transforming your model.", + UserWarning + ) + + transformed_model = transform_model(tce_calculator.models["energy"]) + energy = _as_scalar(transformed_model.predict( + tce_calculator.get_feature_vector(initial_configuration).reshape(1, -1) + )) + return transformed_model, energy def monte_carlo( @@ -203,52 +259,33 @@ def mc_step(atoms: Atoms) -> Atoms: Defaults to `False` for backwards compatibility. """ - if not generator: - generator = np.random.default_rng(seed=0) + generator = np.random.default_rng(seed=0) if generator is None else generator + callback = callback or _default_mc_callback + mc_step = mc_step or _default_mc_step(generator) + energy_modifier = energy_modifier or (lambda initial, final: 0.0) + beta_values = _resolve_beta_values(beta, num_steps) - if not callback: - def callback(step_: int, num_steps_: int): - LOGGER.info(f"MC step {step_:.0f}/{num_steps_:.0f}") + transformed_model, energy = _prepare_energy_model(tce_calculator, initial_configuration) - if not mc_step: - def mc_step(atoms: Atoms) -> Atoms: - new_atoms = atoms.copy() - i, j = generator.integers(len(atoms), size=2) - new_atoms[i].symbol, new_atoms[j].symbol = new_atoms[j].symbol, new_atoms[i].symbol - return new_atoms + def _advance_mc_step(initial_configuration: Atoms, energy: float, step: int) -> tuple[Atoms, float]: + """Apply one Metropolis step and return the updated state.""" - if not energy_modifier: - def energy_modifier(initial: Atoms, final: Atoms) -> float: - return 0.0 - - if isinstance(beta, (Sequence, np.ndarray)): - assert len(beta) == num_steps, "if beta is a sequence, it must be the same length as num_steps" - beta_values = np.array(beta) - elif isinstance(beta, float): - beta_values = np.full(num_steps, beta) - else: - raise TypeError("beta must be either a float or a sequence of floats") + new_configuration = mc_step(initial_configuration) + feature_diff = tce_calculator.get_feature_vector_difference( + initial_configuration, new_configuration + ).reshape(1, -1) + energy_diff = transformed_model.predict(feature_diff) + energy_diff += energy_modifier(initial_configuration, new_configuration) + if not isinstance(energy_diff, float): + energy_diff = energy_diff.item() - # try to pass zeros into the model - zero_feature = np.zeros(tce_calculator.feature_vector_size).reshape(1, -1) - predicted = tce_calculator.models["energy"].predict(zero_feature) - if isinstance(predicted, np.ndarray): - predicted = predicted.item() - if predicted != 0.0: - warnings.warn( - "Input model has an intercept, which will mess with energy difference calculations. " - "The monte carlo run will automatically zero-out this intercept, transforming your model.", - UserWarning - ) - - transformed_model = transform_model(tce_calculator.models["energy"]) + if np.exp(-beta_values[step] * energy_diff) > 1.0 - generator.random(): + LOGGER.debug(f"move accepted with energy difference {energy_diff}") + initial_configuration = new_configuration + energy += energy_diff - energy = transformed_model.predict( - tce_calculator.get_feature_vector(initial_configuration).reshape(1, -1) - ) - if isinstance(energy, np.ndarray): - energy = energy.item() + return initial_configuration, energy def _generating_fn(initial_configuration: Atoms, energy: float) -> Generator[Atoms, None, None]: """Wrap the generator logic in a function so that we can return both output types.""" @@ -262,21 +299,7 @@ def _generating_fn(initial_configuration: Atoms, energy: float) -> Generator[Ato yield to_save LOGGER.info(f"saved configuration at step {step:.0f}/{num_steps:.0f}") - new_configuration = mc_step(initial_configuration) - feature_diff = tce_calculator.get_feature_vector_difference( - initial_configuration, new_configuration - ).reshape(1, -1) - energy_diff = transformed_model.predict(feature_diff) - energy_diff += energy_modifier(initial_configuration, new_configuration) - - if not isinstance(energy_diff, float): - energy_diff = energy_diff.item() - if np.exp(-beta_values[step] * energy_diff) > 1.0 - generator.random(): - LOGGER.debug(f"move accepted with energy difference {energy_diff}") - initial_configuration = new_configuration - energy += energy_diff - - if return_generator: - return _generating_fn(initial_configuration, energy) - - return list(_generating_fn(initial_configuration, energy)) + initial_configuration, energy = _advance_mc_step(initial_configuration, energy, step) + + frames = _generating_fn(initial_configuration, energy) + return frames if return_generator else list(frames)