Skip to content

Commit

Permalink
postprocess
Browse files Browse the repository at this point in the history
  • Loading branch information
yomichi committed Aug 21, 2023
1 parent 9db8dd9 commit 37b77b0
Show file tree
Hide file tree
Showing 11 changed files with 900 additions and 253 deletions.
29 changes: 17 additions & 12 deletions abics/applications/latgas_abinitio_interface/default_observer.py
Original file line number Diff line number Diff line change
Expand Up @@ -81,7 +81,7 @@ class DefaultObserver(ObserverBase):
reference_species: list[str]
calculators: list

def __init__(self, comm, Lreload=False, params={}):
def __init__(self, comm, Lreload=False, params={}, with_energy=True):
"""
Parameters
Expand All @@ -100,7 +100,11 @@ def __init__(self, comm, Lreload=False, params={}):
with open(os.path.join(str(myrank), "obs.dat"), "r") as f:
self.lprintcount = int(f.readlines()[-1].split()[0]) + 1

self.names = ["energy_internal", "energy"]
self.with_energy = with_energy
if with_energy:
self.names = ["energy_internal", "energy"]
else:
self.names = []

params_solvers = params.get("solver", [])
self.calculators = []
Expand Down Expand Up @@ -162,16 +166,17 @@ def logfunc(self, calc_state: MCAlgorithm) -> tuple[float, ...]:
assert calc_state.config is not None
structure: Structure = calc_state.config.structure_norel

energy_internal = calc_state.model.internal_energy(calc_state.config)
energy = calc_state.model.energy(calc_state.config)

result = [energy_internal, energy]

if self.minE is None or energy < self.minE:
self.minE = energy
with open("minEfi.dat", "a") as f:
f.write(str(self.minE) + "\n")
structure.to(fmt="POSCAR", filename="minE.vasp")
if self.with_energy:
energy_internal = calc_state.model.internal_energy(calc_state.config)
energy = calc_state.model.energy(calc_state.config)
result = [energy_internal, energy]
if self.minE is None or energy < self.minE:
self.minE = energy
with open("minEfi.dat", "a") as f:
f.write(str(self.minE) + "\n")
structure.to(fmt="POSCAR", filename="minE.vasp")
else:
result = []

for calculator in self.calculators:
name = calculator["name"]
Expand Down
3 changes: 2 additions & 1 deletion abics/applications/latgas_abinitio_interface/model_setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -1317,7 +1317,7 @@ def reset_from_structure(self, st_in):
st = map2perflat(st, st_in)
st.remove_species(["X"])

for defect_sublattice in self.defect_sublattices:
for defect_sublattice in self.defect_sublattices:
# Extract group information for this sublattice
d_sp2grp = {}
sp = set()
Expand Down Expand Up @@ -1359,6 +1359,7 @@ def reset_from_structure(self, st_in):
raise InputError(
"initial structure violates constraints specified by constraint_func"
)
self.structure_norel = copy.deepcopy(self.structure)

def dummy_structure(self):
"""
Expand Down
214 changes: 121 additions & 93 deletions abics/sampling/pamc.py
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@

logger = logging.getLogger("main")


class PAMCParams:
"""Parameter set for population annealing Monte Carlo
Expand Down Expand Up @@ -193,7 +194,7 @@ def __init__(
self.use_resample_old = False

def reload(self):
# not supported yet
# not supported yet
pass
# self.mycalc.config = pickle_load(os.path.join(str(self.rank), "calc.pickle"))
# self.obs_save0 = numpy_load(os.path.join(str(self.rank), "obs_save.npy"))
Expand All @@ -220,6 +221,7 @@ def save(self, save_obs: bool):
obs_save_ = np.array(self.obs_save)
numpy_save(obs_save_, "obs_save.npy")
numpy_save(self.logweight_history, "logweight_hist.npy")
numpy_save(self.dlogz, "dlogz.npy")
numpy_save(self.kT_history, "kT_hist.npy")

def anneal(self, energy: float):
Expand Down Expand Up @@ -330,9 +332,9 @@ def run(
kTnum = len(self.kTs)
nba = nsteps // kTnum
self.nsteps_between_anneal = nba * np.ones(kTnum, dtype=int)
if nsteps - nba*kTnum > 0:
self.nsteps_between_anneal[-(nsteps - nba*kTnum):] += 1
self.isamples_offsets = np.zeros(kTnum+1, dtype=int)
if nsteps - nba * kTnum > 0:
self.nsteps_between_anneal[-(nsteps - nba * kTnum) :] += 1
self.isamples_offsets = np.zeros(kTnum + 1, dtype=int)

if subdirs:
try:
Expand All @@ -341,7 +343,7 @@ def run(
pass
os.chdir(str(self.rank))
if not self.Lreload:
#self.mycalc.config.shuffle()
# self.mycalc.config.shuffle()
self.mycalc.energy = self.mycalc.model.energy(self.mycalc.config)
self.Tindex = 0
self.anneal(self.mycalc.energy)
Expand All @@ -356,7 +358,11 @@ def run(
while self.Tindex < numT:
if self.Tindex > 0:
self.anneal(self.mycalc.energy)
logger.info("--Anneal from T={} to {}".format(self.kTs[self.Tindex - 1], self.kTs[self.Tindex]))
logger.info(
"--Anneal from T={} to {}".format(
self.kTs[self.Tindex - 1], self.kTs[self.Tindex]
)
)
if self.Tindex % resample_frequency == 0:
self.resample()
logger.info("--Resampling finishes")
Expand Down Expand Up @@ -407,93 +413,115 @@ def run(
os.chdir("../")

def postproc(self):
nT = len(self.kTs)
obs_arr = np.array(self.obs_save)
nobs = obs_arr.shape[1]

obs1 = np.zeros((nT, nobs))
obs2 = np.zeros((nT, nobs))
logweights = np.zeros(nT)
dlogz = np.array(self.dlogz)

for i, (offset, offset2) in enumerate(zip(self.isamples_offsets[:-1],
self.isamples_offsets[1:])):
logweights[i] = self.logweight_history[offset]
o_a = obs_arr[offset:offset2, :]
obs1[i, :] += o_a.mean(axis=0)
obs2[i, :] += (o_a*o_a).mean(axis=0)

obs = np.zeros((nT, 2 * nobs), dtype=np.float64)
for iobs in range(nobs):
obs[:, 2 * iobs] = obs1[:, iobs]
obs[:, 2 * iobs + 1] = obs2[:, iobs]

nreplicas = self.procs
buffer_all = np.zeros((nreplicas, nT, 2 + 2 * nobs), dtype=np.float64)
self.comm.Allgather(
np.concatenate(
[logweights.reshape(-1, 1), dlogz.reshape(-1, 1), obs], axis=1
),
buffer_all,
kTs = self.kTs
obs_save = self.obs_save
dlogz = self.dlogz
logweight_history = self.logweight_history
obsnames = self.obsnames
isamples_offsets = self.isamples_offsets
comm = self.comm
postproc(
kTs, obs_save, dlogz, logweight_history, obsnames, isamples_offsets, comm
)

logweights_all = buffer_all[:, :, 0]
dlogz_all = buffer_all[:, :, 1]
obs_all = buffer_all[:, :, 2:]

logweights_max = logweights_all.max(axis=0)
weights = np.exp(logweights_all - logweights_max)
ow = np.einsum("ito,it->ito", obs_all, weights)

lzw = logweights_all + dlogz_all - logweights_max
lzw_max = lzw.max(axis=0)
zw = np.exp(lzw - lzw_max)

# bootstrap
index = np.random.randint(nreplicas, size=nreplicas)
numer = ow[index, :, :].mean(axis=0)
zw_numer = zw[index, :].mean(axis=0)
denom = weights[index, :].mean(axis=0)
o = np.zeros((nT, 3 * nobs + 1), dtype=np.float64)
for iobs in range(nobs):
o[:, 3 * iobs] = numer[:, 2 * iobs] / denom[:]
o[:, 3 * iobs + 1] = numer[:, 2 * iobs + 1] / denom[:]
o[:, 3 * iobs + 2] = o[:, 3 * iobs + 1] - o[:, 3 * iobs] ** 2
o[:, 3 * nobs] = zw_numer / denom
o_all = np.zeros((nreplicas, nT, 3 * nobs + 1), dtype=np.float64)
self.comm.Allgather(o, o_all)
if self.rank == 0:
o_mean = o_all.mean(axis=0)
o_err = o_all.std(axis=0)
for iobs, oname in enumerate(self.obsnames):
with open(f"{oname}.dat", "w") as f:
f.write( "# $1: temperature\n")
f.write(f"# $2: <{oname}>\n")
f.write(f"# $3: ERROR of <{oname}>\n")
f.write(f"# $4: <{oname}^2>\n")
f.write(f"# $5: ERROR of <{oname}^2>\n")
f.write(f"# $6: <{oname}^2> - <{oname}>^2\n")
f.write(f"# $7: ERROR of <{oname}^2> - <{oname}>^2\n")

for iT in range(nT):
f.write(f"{self.kTs[iT]}")
for j in range(3):
f.write(f" {o_mean[iT, 3*iobs+j]} {o_err[iT, 3*iobs+j]}")
f.write("\n")
dlogZ = np.log(o_mean[:, 3 * nobs]) + lzw_max[:]
dlogZ_err = o_err[:, 3 * nobs] / o_mean[:, 3 * nobs]
with open("logZ.dat", "w") as f:
f.write("# $1: temperature\n")
f.write("# $2: logZ\n")
f.write("# $3: ERROR of log(Z)\n")
f.write("# $4: log(Z/Z')\n")
f.write("# $5: ERROR of log(Z/Z')\n")
F = 0.0
dF = 0.0
f.write(f"inf 0.0 0.0 0.0 0.0\n")
for i, (dlz, dlz_e) in enumerate(zip(dlogZ, dlogZ_err)):
F += dlz
dF += dlz_e
f.write(f"{self.kTs[i]} {F} {dF} {dlz} {dlz_e}\n")

self.comm.Barrier()
def postproc(
kTs,
obs_save,
dlogz,
logweight_history,
obsnames,
isamples_offsets,
comm,
):
procs = comm.Get_size()
rank = comm.Get_rank()

nT = len(kTs)
obs_arr = np.array(obs_save)
nobs = obs_arr.shape[1]

obs1 = np.zeros((nT, nobs))
obs2 = np.zeros((nT, nobs))
logweights = np.zeros(nT)
dlogz = np.array(dlogz)

for i, (offset, offset2) in enumerate(
zip(isamples_offsets[:-1], isamples_offsets[1:])
):
logweights[i] = logweight_history[offset]
o_a = obs_arr[offset:offset2, :]
obs1[i, :] += o_a.mean(axis=0)
obs2[i, :] += (o_a * o_a).mean(axis=0)

obs = np.zeros((nT, 2 * nobs), dtype=np.float64)
for iobs in range(nobs):
obs[:, 2 * iobs] = obs1[:, iobs]
obs[:, 2 * iobs + 1] = obs2[:, iobs]

nreplicas = procs
buffer_all = np.zeros((nreplicas, nT, 2 + 2 * nobs), dtype=np.float64)
comm.Allgather(
np.concatenate([logweights.reshape(-1, 1), dlogz.reshape(-1, 1), obs], axis=1),
buffer_all,
)

logweights_all = buffer_all[:, :, 0]
dlogz_all = buffer_all[:, :, 1]
obs_all = buffer_all[:, :, 2:]

logweights_max = logweights_all.max(axis=0)
weights = np.exp(logweights_all - logweights_max)
ow = np.einsum("ito,it->ito", obs_all, weights)

lzw = logweights_all + dlogz_all - logweights_max
lzw_max = lzw.max(axis=0)
zw = np.exp(lzw - lzw_max)

# bootstrap
index = np.random.randint(nreplicas, size=nreplicas)
numer = ow[index, :, :].mean(axis=0)
zw_numer = zw[index, :].mean(axis=0)
denom = weights[index, :].mean(axis=0)
o = np.zeros((nT, 3 * nobs + 1), dtype=np.float64)
for iobs in range(nobs):
o[:, 3 * iobs] = numer[:, 2 * iobs] / denom[:]
o[:, 3 * iobs + 1] = numer[:, 2 * iobs + 1] / denom[:]
o[:, 3 * iobs + 2] = o[:, 3 * iobs + 1] - o[:, 3 * iobs] ** 2
o[:, 3 * nobs] = zw_numer / denom
o_all = np.zeros((nreplicas, nT, 3 * nobs + 1), dtype=np.float64)
comm.Allgather(o, o_all)
if rank == 0:
o_mean = o_all.mean(axis=0)
o_err = o_all.std(axis=0)
for iobs, oname in enumerate(obsnames):
with open(f"{oname}.dat", "w") as f:
f.write("# $1: temperature\n")
f.write(f"# $2: <{oname}>\n")
f.write(f"# $3: ERROR of <{oname}>\n")
f.write(f"# $4: <{oname}^2>\n")
f.write(f"# $5: ERROR of <{oname}^2>\n")
f.write(f"# $6: <{oname}^2> - <{oname}>^2\n")
f.write(f"# $7: ERROR of <{oname}^2> - <{oname}>^2\n")

for iT in range(nT):
f.write(f"{kTs[iT]}")
for j in range(3):
f.write(f" {o_mean[iT, 3*iobs+j]} {o_err[iT, 3*iobs+j]}")
f.write("\n")
dlogZ = np.log(o_mean[:, 3 * nobs]) + lzw_max[:]
dlogZ_err = o_err[:, 3 * nobs] / o_mean[:, 3 * nobs]
with open("logZ.dat", "w") as f:
f.write("# $1: temperature\n")
f.write("# $2: logZ\n")
f.write("# $3: ERROR of log(Z)\n")
f.write("# $4: log(Z/Z')\n")
f.write("# $5: ERROR of log(Z/Z')\n")
F = 0.0
dF = 0.0
f.write(f"inf 0.0 0.0 0.0 0.0\n")
for i, (dlz, dlz_e) in enumerate(zip(dlogZ, dlogZ_err)):
F += dlz
dF += dlz_e
f.write(f"{kTs[i]} {F} {dF} {dlz} {dlz_e}\n")
comm.Barrier()
Loading

0 comments on commit 37b77b0

Please sign in to comment.