Clexulator writing updates: write cluster_functions.json.gz and verbo…

…se output

CHANGELOG.md

@@ -13,11 +13,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Added `ClexulatorWriter.generated_files` attribute, holding a list of the files generated when writing a Clexulator
 - Added generation of latex formulas during the Clexulator writing process
 - Added `casm.bset.cluster_functions.MakeVariableName` for variable naming
+- Added `verbose` and `very_verbose` options to `casm.bset.write_clexulator`
+- Added `to_dict` methods to `ClusterFunctionsBuilder`, `OrbitMatrixRepBuilder` and `ClusterMatrixRepBuilder`
+- Added `to_dict` and `from_dict` methods to `ExponentSumConstraint`
+- Added `casm.bset.json_io` module
 
 ### Changed
 
-- Updated `ClexulatorWriter.write` to write a "variables.json" file for each Clexulator (including local Clexulator) which contains the variables used by the jinja2 templates as well as information like basis function formulas generated during the write process
-- Changed the name used for occupation variables. The name "\\phi" is the base, but the number and meaning of the indices can vary depending on the site basis functions. An "occ_site_functions_info" dict in the "variables.json" file contains the string value used as a template ("occ_var_name") and a description of the variable including its indices, if any ("occ_var_desc").
+- Updated `ClexulatorWriter.write` to write a "variables.json.gz" file for each Clexulator (including local Clexulator) which contains the variables used by the jinja2 templates as well as information like basis function formulas generated during the write process
+- Updated `ClexulatorWriter.write` to write a "cluster_functions.json.gz" file with the generated clusters, matrix reps, and functions.
+- Changed the name used for occupation variables. The name "\\phi" is the base, but the number and meaning of the indices can vary depending on the site basis functions. An "occ_site_functions_info" dict in the "variables.json" file contains the string value used as a template ("occ_var_name") and a description of the variable including its indices, if any ("occ_var_desc" and "occ_var_indices").
 
 
 ## [2.0a1] - 2024-08-15

casm/bset/_methods.py

@@ -280,6 +280,8 @@ def write_clexulator(
     version: str = "v1.basic",
     linear_function_indices: Optional[set[int]] = None,
     cpp_fmt: Optional[CppFormatProperties] = None,
+    verbose: bool = True,
+    very_verbose: bool = False,
 ) -> tuple[pathlib.Path, Optional[list[pathlib.Path]], PrimNeighborList]:
     """Write a CASM Clexulator
 
@@ -339,6 +341,12 @@ def write_clexulator(
     cpp_fmt: Optional[CppFormatProperties] = None
         C++ string formatting properties. If None, default constructor values are used.
 
+    verbose: bool = True
+        Print progress statements
+
+    very_verbose: bool = False
+        Print detailed progress statements from the cluster functions builder.
+
     Returns
     -------
     src_path: pathlib.Path
@@ -368,6 +376,8 @@ def write_clexulator(
         prim=prim,
         clex_basis_specs=clex_basis_specs,
         prim_neighbor_list=prim_neighbor_list,
+        verbose=verbose,
+        very_verbose=very_verbose,
     )
 
     return (writer.src_path, writer.local_src_path, prim_neighbor_list)

casm/bset/clexwriter/_methods.py

@@ -1,5 +1,7 @@
 import pathlib
 import re
+import sys
+import time
 from typing import Optional
 
 import jinja2
@@ -21,6 +23,7 @@
     make_local_point_functions,
     make_point_functions,
 )
+from casm.bset.json_io import dump
 from casm.bset.polynomial_functions import (
     PolynomialFunction,
 )
@@ -836,6 +839,8 @@ def write(
         prim: casmconfig.Prim,
         clex_basis_specs: ClexBasisSpecs,
         prim_neighbor_list: PrimNeighborList,
+        verbose: bool = True,
+        very_verbose: bool = False,
     ):
         """Write Clexulator source files and related files
 
@@ -852,6 +857,12 @@ def write(
             variables included in the cluster functions, relative to a reference unit
             cell. If None, a default neighbor list is constructed.
 
+        verbose: bool = True
+            Print progress statements
+
+        very_verbose: bool = False
+            Print detailed progress statements from the cluster functions builder.
+
         """
         cluster_specs = clex_basis_specs.cluster_specs
         bfunc_specs = clex_basis_specs.basis_function_specs
@@ -873,6 +884,14 @@ def write(
                     "site_basis_functions"
                 ]
 
+        # store a list of generated files
+        generated_files = []
+
+        if verbose:
+            start = time.time()
+            print("Building cluster functions...")
+            sys.stdout.flush()
+
         builder = ClusterFunctionsBuilder(
             prim=prim,
             generating_group=cluster_specs.generating_group(),
@@ -883,8 +902,22 @@ def write(
             orbit_branch_max_poly_order=orbit_branch_max_poly_order,
             occ_site_basis_functions_specs=occ_site_basis_functions_specs,
             prim_neighbor_list=prim_neighbor_list,
+            verbose=very_verbose,
         )
 
+        path = self.bset_dir / "cluster_functions.json.gz"
+        generated_files.append(path)
+        dump(builder.to_dict(), path, force=True, gz=True)
+
+        if verbose:
+            print("Building cluster functions DONE")
+            print(f"- n_orbits: {len(builder.clusters)}")
+            print(f"- n_functions: {builder.n_functions}")
+            elapsed_time = time.time() - start
+            print(f"build time: {elapsed_time:0.4f} (s)")
+            print()
+            sys.stdout.flush()
+
         env = jinja2.Environment(
             trim_blocks=True,
             lstrip_blocks=True,
@@ -899,21 +932,32 @@ def write(
             is_periodic=is_periodic,
             prim_neighbor_list=builder.prim_neighbor_list,
             occ_site_functions=builder.occ_site_functions,
+            occ_site_functions_info=builder.occ_site_functions_info,
             linear_function_indices=self.linear_function_indices,
             cpp_fmt=self.cpp_fmt,
+            verbose=verbose,
         )
 
-        # store a list of generated files
-        generated_files = []
-
         ## write periodic clexulator / prototype local clexulator
+        if verbose:
+            start = time.time()
+            print("Generating variables...")
+            sys.stdout.flush()
+
         writer = self.writer_type(
             i_clex=None,
             clusters=builder.clusters,
             functions=builder.functions,
             **writer_params,
         )
 
+        if verbose:
+            print("Generating variables DONE")
+            elapsed_time = time.time() - start
+            print(f"generation time: {elapsed_time:0.4f} (s)")
+            print()
+            sys.stdout.flush()
+
         # Clexulator source file
         path = self.bset_dir / f"{writer.clexulator_name}.cc"
         self.src_path = path
@@ -924,37 +968,45 @@ def write(
         # basis.json file
         path = self.bset_dir / "basis.json"
         generated_files.append(path)
-        with open(path, "w") as f:
-            data = builder.basis_dict(
-                clex_basis_specs=clex_basis_specs,
-                coordinate_mode="frac",
-            )
-            f.write(xtal.pretty_json(data))
+        data = builder.basis_dict(
+            clex_basis_specs=clex_basis_specs, coordinate_mode="frac"
+        )
+        dump(data, path, force=True)
 
         # write all template variables
-        path = self.bset_dir / "variables.json"
+        path = self.bset_dir / "variables.json.gz"
         generated_files.append(path)
-        with open(path, "w") as f:
-            f.write(xtal.pretty_json(writer.variables()))
+        dump(writer.variables(), path, force=True, gz=True)
 
         ## if local clexulators
         if phenomenal is not None:
             # write equivalents_info.json
             path = self.bset_dir / "equivalents_info.json"
             generated_files.append(path)
-            with open(path, "w") as f:
-                f.write(xtal.pretty_json(builder.equivalents_info_dict()))
+            dump(builder.equivalents_info_dict(), path, force=True)
 
             # write each local clexulator
             self.local_src_path = []
             for i_clex in range(len(builder.equivalent_clusters)):
+                if verbose:
+                    start = time.time()
+                    print(f"Generating variables for local clexulator {i_clex}...")
+                    sys.stdout.flush()
+
                 writer = self.writer_type(
                     i_clex=i_clex,
                     clusters=builder.equivalent_clusters[i_clex],
                     functions=builder.equivalent_functions[i_clex],
                     **writer_params,
                 )
 
+                if verbose:
+                    print("Generating variables DONE")
+                    elapsed_time = time.time() - start
+                    print(f"generation time: {elapsed_time:0.4f} (s)")
+                    print()
+                    sys.stdout.flush()
+
                 local_dir = self.bset_dir / f"{i_clex}"
                 local_dir.mkdir(parents=True, exist_ok=True)
 
@@ -968,8 +1020,7 @@ def write(
                 # write all template variables
                 path = local_dir / "variables.json"
                 generated_files.append(path)
-                with open(path, "w") as f:
-                    f.write(xtal.pretty_json(writer.variables()))
+                dump(writer.variables(), path, force=True)
 
         # write generated_files.json (includes itself)
         path = self.bset_dir / "generated_files.json"
@@ -983,6 +1034,5 @@ def write(
             data["local_src_path"] = [
                 str(p.relative_to(self.bset_dir)) for p in self.local_src_path
             ]
-        with open(path, "w") as f:
-            f.write(xtal.pretty_json(data))
+        dump(data, path, force=True)
         return None

casm/bset/clexwriter/_write_v1_basic.py

@@ -1,4 +1,6 @@
 import copy
+import sys
+import time
 from typing import Optional
 
 import libcasm.clusterography as casmclust
@@ -47,8 +49,10 @@ def __init__(
         clusters: list[list[casmclust.Cluster]],
         functions: list[list[list[PolynomialFunction]]],
         occ_site_functions: list[dict],
+        occ_site_functions_info: dict,
         linear_function_indices: Optional[set[int]] = None,
         cpp_fmt: Optional[CppFormatProperties] = None,
+        verbose: bool = False,
     ):
         """
 
@@ -95,6 +99,24 @@ def __init__(
             - `"value"`: list[list[float]], list of the site basis function values,
               as ``value[function_index][occupant_index]``.
 
+        occ_site_functions_info: dict
+            Information about occupation site basis functions.
+
+            Occupation site basis functions info, with format:
+
+            - `"max_function_index"`: int, The maximum site function index, across all
+              sublattices.
+            - `"all_sublattices_have_same_site_functions"`: bool, True if all
+              sublattices have same site functions; False otherwise.
+            - `"occ_var_name"`: str, A variable name template for the site functions,
+              which may be formated using `b` for sublattice index and `m` for site
+              function index (i..e ``occ_var_name.format(b=0, m=1)``).
+            - `"occ_var_desc": str, A description of the occupation
+              variable, including a description of the subscript indices.
+            - `"occ_var_indices"`: list[list[str, str]], A list of lists, where each
+              sublist contains the variable name and description for each subscript
+              index.
+
         linear_function_indices : Optional[set[int]]
             The linear indices of the functions that will be included. If None,
             all functions will be included in the Clexulator. Otherwise,
@@ -111,6 +133,9 @@ def __init__(
                     coeff_atol=1e-10,
                 )
 
+        verbose: bool = False
+            Print progress statements
+
         """
 
         ### defaults ###
@@ -175,6 +200,23 @@ def __init__(
                     "a constant function."
                 )
 
+        self.occ_site_functions_info = occ_site_functions_info
+        """dict: Information about occupation site basis functions.
+        
+        Occupation site basis functions info, with format:
+        
+        - `"max_function_index"`: int, The maximum site function index, across all
+            sublattices.
+        - `"all_sublattices_have_same_site_functions"`: bool, True if all _sublattices
+            have same site functions; False otherwise.
+        - `"occ_var_name"`: str, A variable name template for the site functions,
+            which may be formated using `b` for sublattice index and `m` for site 
+            function index (i..e ``occ_var_name.format(b=0, m=1)``).
+        - `"occ_var_desc": str, A description of the occupation
+            variable, including a description of the subscript indices.
+            
+        """
+
         ## set linear_function_indices
         self.linear_function_indices = linear_function_indices
         """Optional[set[int]]: The linear indices of the functions that will be 
@@ -204,6 +246,11 @@ def __init__(
             self.n_corr += len(prototype_functions)
 
         ## set neighborhood info
+        if verbose:
+            start = time.time()
+            print("Generating neighborhoods...")
+            sys.stdout.flush()
+
         _complete_neighborhood, _function_neighborhoods = make_neighborhoods(
             is_periodic=is_periodic,
             prim_neighbor_list=prim_neighbor_list,
@@ -212,6 +259,13 @@ def __init__(
             linear_function_indices=linear_function_indices,
         )
 
+        if verbose:
+            print("Generating neighborhoods DONE")
+            elapsed_time = time.time() - start
+            print(f"time: {elapsed_time:0.4f} (s)")
+            print()
+            sys.stdout.flush()
+
         self.complete_neighborhood = _complete_neighborhood
         """dict: The neighborhood needed to evaluate basis functions 
         involving sites in the origin unit cell. 
@@ -281,6 +335,11 @@ def __init__(
         continuous DoF are None.
         """
 
+        if verbose:
+            start = time.time()
+            print("Generating orbit function formulas...")
+            sys.stdout.flush()
+
         self.orbit_bfuncs, self.orbit_bfuncs_variables_needed = make_orbit_bfuncs(
             prim_neighbor_list=self.prim_neighbor_list,
             clusters=clusters,
@@ -290,6 +349,13 @@ def __init__(
             linear_function_indices=self.linear_function_indices,
         )
 
+        if verbose:
+            print("Generating orbit function formulas DONE")
+            elapsed_time = time.time() - start
+            print(f"time: {elapsed_time:0.4f} (s)")
+            print()
+            sys.stdout.flush()
+
         ## set site_bfuncs and site_bfuncs_variables_needed_at
 
         self.site_bfuncs = None
@@ -326,6 +392,11 @@ def __init__(
         continuous DoF are None.
         """
 
+        if verbose:
+            start = time.time()
+            print("Generating site function formulas...")
+            sys.stdout.flush()
+
         self.site_bfuncs, self.site_bfuncs_variables_needed_at = make_site_bfuncs(
             is_periodic=is_periodic,
             prim_neighbor_list=self.prim_neighbor_list,
@@ -336,6 +407,13 @@ def __init__(
             linear_function_indices=self.linear_function_indices,
         )
 
+        if verbose:
+            print("Generating site function formulas DONE")
+            elapsed_time = time.time() - start
+            print(f"time: {elapsed_time:0.4f} (s)")
+            print()
+            sys.stdout.flush()
+
         ## set n_point_corr_sites
         n_point_corr_sites = 0
         for x in self.site_bfuncs:
@@ -500,6 +578,7 @@ def variables(self) -> dict:
             "site_bfuncs": self.site_bfuncs,
             "site_bfuncs_variables_needed_at": self.site_bfuncs_variables_needed_at,
             "occ_site_functions": self.occ_site_functions,
+            "occ_site_functions_info": self.occ_site_functions_info,
             "continuous_dof": self.continuous_dof,
             "params": self.params,
             "nlist_weight_matrix": self.prim_neighbor_list.weight_matrix().tolist(),