Add docstrings

burnman/data/input_raw_endmember_datasets/HeFESTo_2024_to_burnman.py

@@ -23,6 +23,7 @@
 from burnman.constants import gas_constant
 from itertools import groupby
 import re
+from datetime import date
 
 
 def all_equal(iterable):
@@ -40,7 +41,7 @@ def rfloat(x, m=1.0):
 Crystallography of wustites shows that the locations of the tetrahedral site
 and the octahedral vacancies are highly correlated, with the tetrahedral
 site being surrounded by a tetrahedral array of octahedral vacancies
-(Gavarri et al. 1979 ).Because of this highly structured arrangement,
+(Gavarri et al. 1979). Because of this highly structured arrangement,
 we ignore the contribution to the ideal entropy of mixing from the
 tetrahedral site and the vacant octahedral sites.
 
@@ -81,6 +82,8 @@ def rfloat(x, m=1.0):
     "hmag": {"p": 0.4},
 }
 
+solution_dataset = "SLB_2024"
+paper_reference = "Stixrude and Lithgow-Bertelloni (2024)"
 hefesto_path = "HeFESTo_parameters_010123"
 mbrdir = pathlib.Path(hefesto_path)
 phasedir = pathlib.Path(f"{hefesto_path}/phase")
@@ -409,19 +412,42 @@ def rfloat(x, m=1.0):
 
 # WRITE FILE
 
-print(
-    "# This file is part of BurnMan - a thermoelastic and thermodynamic toolkit\n"
-    "# for the Earth and Planetary Sciences\n"
-    "# Copyright (C) 2012 - 2024 by the BurnMan team,\n"
-    "# released under the GNU GPL v2 or later.\n"
-    "\n\n"
+current_year = str(date.today())[:4]
+underline = "^" * len(solution_dataset)
+preamble = (
+    "# This file is part of BurnMan - a thermoelastic\n"
+    "# and thermodynamic toolkit for the Earth and Planetary Sciences\n"
+    f"# Copyright (C) 2012 - {current_year} by the BurnMan team, released under the GNU\n"
+    "# GPL v2 or later.\n"
+    "\n"
     '"""\n'
-    "SLB_2024\n"
-    "Minerals from Stixrude & Lithgow-Bertelloni 2024 and references therein\n"
-    f"File autogenerated using HeFESTO_2024_to_burnman.py and {hefesto_path}\n"
+    f"{solution_dataset}\n"
+    f"{underline}\n\n"
+    f"Minerals from {paper_reference}, and references therein.\n\n"
+    "Contains the following solutions:"
+)
+
+print(preamble)
+for k, v in sorted(sol_params.items()):
+    print(f"* {solution_aliases[k]} ({k})")
+
+print(
+    "\nRelaxed solutions (which account for softening through the spin transition):\n"
+)
+print("* ferropericlase_relaxed")
+print("* bridgmanite_relaxed")
+
+print("\nEndmember phases:")
+for k, v in sorted(mbr_names.items()):
+    print(f"* {v} ({k})")
+print()
+
+print(
+    f"File autogenerated using HeFESTO_2024_to_burnman.py and {hefesto_path}.\n"
     '"""\n\n'
     "\n"
     "\n\n"
+    "import numpy as np\n"
     "from ..eos.slb import SLB3\n"
     "from ..eos import debye\n"
     "from ..eos import birch_murnaghan as bm\n"

burnman/data/input_raw_endmember_datasets/hpx_eos_to_burnman.py

@@ -179,9 +179,11 @@ def replace_numbers(match):
     n_data = len(data)
 
     noods = []
+    names = []
     for ind in start_indices:
         i = ind
         name = data[i][0]
+        names.append(name)
         n_mbrs = int(data[i][1])
         logging.debug(f"Solution: {name}")
         logging.debug(f"Number of endmembers: {n_mbrs}")
@@ -380,8 +382,24 @@ def replace_numbers(match):
 
             mbr_initializations.append(f'[{mbr}, "{site_formula}"]')
 
+        # Create docstring
+        doc = '        """\n'
+        doc += f"        Initialisation for a {name} solution object.\n"
+
+        doc += "        Contains the following endmembers with associated site occupancies:\n"
+
+        for mbr_initialization in mbr_initializations:
+            doc += f"        * {mbr_initialization}\n"
+
+        if formulation == "asf":
+            doc += "\n        This is implemented as an asymmetric solution.\n"
+        else:
+            doc += "\n        This is implemented as a symmetric solution.\n"
+        doc += '        """\n'
+
         out_ss += f"class {name}(Solution):\n"
         out_ss += "    def __init__(self, molar_fractions=None):\n"
+        out_ss += doc
         out_ss += f'        self.name = "{name}"\n'
         if formulation == "asf":
             out_ss += "        self.solution_model = AsymmetricRegularSolution(\n"
@@ -427,9 +445,15 @@ def replace_numbers(match):
         f"{underline}\n"
         "\n"
         "HPx-eos solutions using endmembers from\n"
-        f"dataset {mbr_dataset}.\n"
-        "The values in this document are all in S.I. units,\n"
-        "unlike those in the original THERMOCALC file.\n"
+        f"dataset {mbr_dataset}.\n\n"
+        "Contains the following solutions:\n"
+    )
+    for name in names:
+        preamble += f"* {name}\n"
+
+    preamble += (
+        "\nThe values in this document are all in S.I. units,\n"
+        "unlike those in the original THERMOCALC file.\n\n"
         "This file is autogenerated using hpx_eos_to_burnman.py\n"
         '"""\n\n'
         f"import numpy as np\n"