Allow for single walker runs, and fix underflow in analysis

pymatnext/__version__.py

@@ -1 +1 @@
-__version__ = "0.1.3"
+__version__ = "0.1.4"

pymatnext/analysis/utils.py

@@ -44,7 +44,7 @@ def calc_log_a(iters, n_walkers, n_cull, discrete=False):
     return log_a
 
 
-def calc_Z_terms(beta, log_a, Es, flat_V_prior=False, N_atoms=None, Vs=None):
+def calc_log_Z_terms(beta, log_a, Es, flat_V_prior=False, N_atoms=None, Vs=None):
     """Return the terms that sum to Z
 
     Parameters
@@ -64,7 +64,7 @@ def calc_Z_terms(beta, log_a, Es, flat_V_prior=False, N_atoms=None, Vs=None):
 
     Returns
     -------
-    Z_term: list of terms that sum to Z, multipled by exp(-shift)
+    log_Z_term: list of logs of terms that sum to Z, shifted by -log_shift
     log_shift: shift subtracted from each log(Z_term_true) to get log(Z_term)
     """
     log_Z_term = log_a[:] - beta * Es[:]
@@ -75,9 +75,10 @@ def calc_Z_terms(beta, log_a, Es, flat_V_prior=False, N_atoms=None, Vs=None):
         log_Z_term += N_atoms * np.log(Vs[:])
 
     log_shift = np.amax(log_Z_term[:])
-    Z_term = np.exp(log_Z_term[:] - log_shift)
+    # Z_term = np.exp(log_Z_term[:] - log_shift)
+    log_Z_term -= log_shift
 
-    return (Z_term, log_shift)
+    return (log_Z_term, log_shift)
 
 
 def analyse_T(T, Es, E_shift, Vs, extra_vals, log_a, flat_V_prior, N_atoms, kB, n_extra_DOF, p_entropy_min=5.0, sum_f=np.sum):
@@ -113,6 +114,8 @@ def analyse_T(T, Es, E_shift, Vs, extra_vals, log_a, flat_V_prior, N_atoms, kB,
         added analytically to energies and specific heats
     p_entropy_min: float, default 5
         minimum value of entropy of probability distribution that indicates a problem (poor sampling, e.g. with P reweighting)
+    sum_f: function
+        function with API equivalent to np.sum, e.g. more accurate sum
 
     Returns
     -------
@@ -121,7 +124,8 @@ def analyse_T(T, Es, E_shift, Vs, extra_vals, log_a, flat_V_prior, N_atoms, kB,
     beta = 1.0 / (kB * T)
 
     # Z_term here is actually Z_term_true * exp(-log_shift)
-    (Z_term, log_shift) = calc_Z_terms(beta, log_a, Es, flat_V_prior, N_atoms, Vs)
+    (log_Z_term, log_shift) = calc_log_Z_terms(beta, log_a, Es, flat_V_prior, N_atoms, Vs)
+    Z_term = np.exp(log_Z_term)
 
     # Note that
     #     Z_term = Z_term_true * exp(-log_shift)
@@ -133,9 +137,9 @@ def analyse_T(T, Es, E_shift, Vs, extra_vals, log_a, flat_V_prior, N_atoms, kB,
 
     if N_atoms is not None:
         N = sum_f(Z_term * N_atoms) / Z_term_sum
-        n_extra_DOF * N
 
-    U = n_extra_DOF / (2.0 * beta) + U_pot + E_shift
+    U_extra_DOF = n_extra_DOF / (2.0 * beta)
+    U = U_pot + U_extra_DOF + E_shift
 
     Cvp = n_extra_DOF * kB / 2.0 + kB * beta * beta * (sum(Z_term * Es**2) / Z_term_sum - U_pot**2)
 
@@ -154,25 +158,27 @@ def analyse_T(T, Es, E_shift, Vs, extra_vals, log_a, flat_V_prior, N_atoms, kB,
     # undo shift of Z_term
     log_Z = np.log(Z_term_sum) + log_shift
 
+    # to make sure that Helmholtz_F approaches U at T -> 0,
     # we want last Z term to have w = 1, so we define a factor f which scales it correctly
     #    f exp(log_shift) Z_term[-1] = 1.0 * exp(-beta Es[-1])
     #    f = exp(-beta Es[-1] - log_shift) / Z_term[-1]
     #    log(f) = -beta Es[-1] - log_shift - log(Z_term[-1])
-    log_f = -beta * Es[-1] - log_shift - np.log(Z_term[-1])
+    log_f = -beta * Es[-1] - log_shift - log_Z_term[-1]
     # this factor rescales every term in Z
     log_Z += log_f
 
     # also add the E_shift
-    Helmholtz_F = -log_Z / beta + E_shift
+    Helmholtz_F = -log_Z / beta + U_extra_DOF + E_shift
 
     mode_config = np.argmax(Z_term)
 
-
     results_dict = {'log_Z': log_Z,
                     'FG': Helmholtz_F,
                     'U': U,
                     'S': (U - Helmholtz_F) * beta,
                     'Cvp': Cvp}
+    if N_atoms is not None:
+        results_dict['N'] = N
 
     if Vs is not None:
         results_dict['V'] = V

pymatnext/cli/ns_analyse.py

@@ -44,8 +44,10 @@ def main():
     pressure_g.add_argument('--delta_P',  help="""delta pressure to use for reweighting (works best with flat V prior)""", type=float)
     p.add_argument('--entropy', '-S', action='store_true', help="""compute and print entropy (relative to entropy of lowest T structure""")
     p.add_argument('--probability_entropy_minimum', type=float, help="""probability entropy mininum that indicates a problem with sampling""", default=5.0)
-    p.add_argument('--plot', '-p', nargs='*', help="""column names to plot, or optionally 'log(colname)'. """
-                                                   """If no column names provided, list allowed names and abort""")
+    p.add_argument('--plot', '-p', nargs='*', help="""column names to plot. Python expression can be used, with '{colname}' """
+                                                   """replaced by the value of the quantity, and 'natoms' by the number of atoms. """
+                                                   """Plotted on a semi-log axis if name or expression is enclosed by 'log(..)'. """
+                                                   """If no column names provided, list allowed names and abort.""")
     p.add_argument('--plot_together', help="""output filename for combined plot""")
     p.add_argument('--plot_together_filenames', action='store_true', help="""show filenames in combined plot""")
     p.add_argument('--plot_twinx_spacing', type=float, help="""spacing for extra twinx y axes""", default=0.15)
@@ -85,13 +87,17 @@ def main():
     ax = {}
 
     def colname(colname_str):
-        m = re.match(r'log\(([^)]*)\)$', colname_str)
+        # strip off optional log()
+        m = re.match(r'^log\((.*)\)$', colname_str)
+        if m:
+            colname_str = m.group(1)
+        # extract col name
+        m = re.match(r'\{([^}]*)\}', colname_str)
         if m:
             return m.group(1)
         else:
             return colname_str
 
-
     for infile_i, infile in enumerate(args.infile):
         iters = []
         Es = []
@@ -346,8 +352,17 @@ def str_format(fmt):
                     fig = Figure()
                     ax = {}
                 for field_i, pfield in enumerate(args.plot):
-                    # should this be done here?  should it be more general, e.g. eval()?
-                    col_log = pfield.startswith('log')
+                    # check for log axis
+                    col_log = re.match(r'^log\(.*\)$', pfield)
+
+                    # check for arb math
+                    m = re.match(r'^(.*){([^}]*)}(.*)$', pfield)
+                    if m:
+                        expr = m.group(1) + m.group(2) + m.group(3)
+                        c = colname(pfield)
+                        plot_data[c] = np.asarray([eval(expr, {}, {'natoms': results_dict['N'], c: v}) for v in plot_data[c]])
+
+                    # replace pfield with colname
                     pfield = colname(pfield)
                     if len(ax) == 0:
                         ax[pfield] = fig.add_subplot()
@@ -391,11 +406,10 @@ def do_plot_sections(pfield, linestyle, color, label):
                                 section_start = T_i
                         plot_section_start = max(section_start - 1, 0)
                         plot_section_end = len(plot_data['T'])
-                        if not got_label:
-                            use_label = label
+
                         pp(plot_data['T'][plot_section_start:plot_section_end], plot_data[pfield][plot_section_start:plot_section_end],
                            linestyle if valid_Ts_bool[section_start] else ':',
-                           color=color, label=use_label)
+                           color=color, label=None if got_label else label)
 
                         ax[pfield].set_ylabel(header_col(pfield))
 

pymatnext/cli/sample.py

@@ -233,6 +233,12 @@ def sample(args, MPI, NS_comm, walker_comm):
         print("params = ", end="")
         pprint.pprint(params, sort_dicts=False)
 
+    # Only for one walker runs, allow culled config to be source of clone
+    if ns.n_configs_global == 1:
+        clone_index_exclude = 0
+    else:
+        clone_index_exclude = 1
+
     time_prev_stdout_report = time.time()
     for loop_iter in loop_iterable:
         if exit_cond(ns, loop_iter):
@@ -254,7 +260,8 @@ def sample(args, MPI, NS_comm, walker_comm):
                               adjust_factor=params_step_size_tune["adjust_factor"])
 
         # pick random config as source for clone.
-        global_ind_of_clone_source = (global_ind_of_max + 1 + ns.rng_global.integers(0, ns.n_configs_global - 1)) % ns.n_configs_global
+        global_ind_of_clone_source = (global_ind_of_max + clone_index_exclude +
+                                      ns.rng_global.integers(0, ns.n_configs_global - clone_index_exclude)) % ns.n_configs_global
         rank_of_clone_source, local_ind_of_clone_source = ns.local_ind(global_ind_of_clone_source)
 
         if clone_history_file: