async def _apply_smoother(self, hess: bool):
import adios2
names = []
if hess and self.precondition:
src = 'hess_vel_raw.bp'
dst = 'hess_vel_smooth.bp'
cmd = f'xconvert_hessian kernels_raw.bp DATABASES_MPI/solver_data.bp hess_vel_raw.bp {self.precondition}'
await self.mpiexec(getpath('adios', 'bin', cmd), getsize(self), 1,
0, 'adios')
else:
src = 'kernels_raw.bp'
dst = 'hessian_smooth.bp' if hess else 'kernels_smooth.bp'
# get the names of the kernels to be smoothed
with adios2.open(self.abs(src), 'r') as fh: # type: ignore
pf = '_crust_mantle/array'
for fstep in fh:
step_vars = fstep.available_variables()
for name in step_vars:
if name.endswith(pf):
name = name.split(pf)[0]
if name.startswith('hess_'):
if hess:
names.append(name)
else:
if not hess:
names.append(name)
# save the number of kernels being smoothed for probe_smoother
kind = 'smooth_' + ('hess' if hess else 'kl')
cache[kind] = len(names)
# get the command to call smoother
cmd = 'bin/xsmooth_laplacian_sem_adios'
radius = self.smooth_hessian if hess else self.smooth_kernels
if isinstance(radius, list):
radius = max(radius[1], radius[0] * radius[2]**self.iteration)
kl = ','.join(names)
await self.mpiexec(
f'{cmd} {radius} {radius} {kl} {src} DATABASES_MPI/ {dst} > OUTPUT_FILES/{kind}.txt',
getsize(self), 1, 0, 'smooth_kernels')
# reset status
del cache[kind]
def add_step(self):
step = len(self.misfits) - 1
self.add(ws := Workspace(f'step_{step:02d}'))
ws.add(ws.mkdir)
# update model
model = self.abs('../model_init.bp')
if step == 0:
mesh = self.abs('../kernel/solver_synthetic/DATABASES_MPI/solver_data.bp')
else:
mesh = self.abs(f'step_{step-1:02d}/kernel_misfit/solver_synthetic/DATABASES_MPI/solver_data.bp')
cmd = f'{self.abs("../../adios/xupdate_model")} {self.steps[-1]} {model} {mesh} {self.abs("../direction.bp")} .'
ws.add(partial(ws.mpiexec, cmd, nprocs=getsize()))
# compute misfit
ws.add(kernel := create_kernel('kernel_misfit', {
'misfit_only': True, 'path_model': ws.abs('model_gll.bp'), 'path_encoded': self.abs('../kernel/observed.ft.h5')
}))
# compute next step
ws.add(partial(self.bracket, kernel))
async def _apply_preconditioner(self):
kl = 'kernels_smooth.bp' if self.smooth_kernels else 'kernels_raw.bp'
hess = 'hess_vel_smooth.bp' if self.smooth_hessian else 'kernels_raw.bp'
cmd = f'xprecond_kernels {kl} {hess} kernels_precond.bp {self.precondition}'
await self.mpiexec(getpath('adios', 'bin', cmd), getsize(self), 1, 0,
'adios')
async def compute_direction(self, ws: Directory, i: int):
if i == 0:
# steepest descent for the first iteration
await super().compute_direction(ws, i)
else:
# input paths
path_file = ws.abs('path.txt')
direction = ws.abs('direction.bp')
solver_file = ws.abs(
'kernel/solver_synthetic/DATABASES_MPI/solver_data.bp')
nprocs = getsize()
# write path to history kernels and directions
lines = [str(min(self.mem, i))]
for j in range(max(0, i - self.mem), i):
lines.append(self.abs(f'iter_{j:02d}/kernels.bp'))
lines.append(self.abs(f'iter_{j:02d}/direction.bp'))
lines.append(ws.abs('kernels.bp'))
ws.write('\n'.join(lines), path_file)
await ws.mpiexec(
f'../adios/xlbfgs {path_file} {solver_file} {direction}',
nprocs=nprocs,
walltime='adios')
async def compute_direction(self, ws: Directory, i: int):
if i == 0:
# steepest descent for the first iteration
await super().compute_direction(ws, i)
else:
# input paths
direction = ws.abs('direction.bp')
kernels = ws.abs(f'kernels.bp')
solver_file = ws.abs('kernel/solver_synthetic/DATABASES_MPI/solver_data.bp')
nprocs = getsize()
# previous iteration
direction0 = self.abs(f'iter_{i-1:02d}/direction.bp')
kernels0 = self.abs(f'iter_{i-1:02d}/kernels.bp')
cmd = f'../adios/xcg_direction {kernels0} {kernels} {direction0} {solver_file} {direction}'
await ws.mpiexec(cmd, nprocs=nprocs, walltime='adios')
def setup(self):
self.clear()
# number of processors to use
nprocs = getsize()
if self.path_forward:
# prepare adjoint simulation
self.add(self._setup_adjoint)
# call solver
self.add(partial(self.mpiexec, 'bin/xspecfem3D', nprocs, 1, 1,
'solver_adjoint', True),
prober=partial(probe_solver, self))
# move OUTPUT_FILES/kernels.bp to kernels_raw.bp
self.add(
partial(self.mv, 'OUTPUT_FILES/kernels.bp', 'kernels_raw.bp'),
'move_kernels')
# smooth kernels / hessian
if self.smooth_kernels:
self.add(partial(self._apply_smoother, False),
'smooth_kernels',
prober=partial(probe_smoother, self, False))
if self.smooth_hessian:
self.add(partial(self._apply_smoother, True),
'smooth_hessian',
prober=partial(probe_smoother, self, True))
# apply preconditioner
if self.precondition:
self.add(self._apply_preconditioner)
# link final kernels to kernels.bp
self.add(self._finalize_adjoint)
else:
# prepare forward simulation
self.add(self._setup_forward)
# call mesher and solver
self.add(partial(self.mpiexec, 'bin/xmeshfem3D', nprocs, 1, 0,
'mesher'),
prober=partial(probe_mesher, self))
self.add(partial(self.mpiexec, 'bin/xspecfem3D', nprocs, 1, 1,
'solver_forward', True),
prober=partial(probe_solver, self))
# move OUTPUT_FILES/synthetic.h5 to traces_raw.h5
self.add(
partial(self.mv, 'OUTPUT_FILES/synthetic.h5', 'traces_raw.h5'),
'move_traces')
# process traces
if self.process_traces:
if 'dst' in self.process_traces:
self.add(
asdf_task(self.abs('traces_raw.h5'),
**self.process_traces))
else:
self.add(
asdf_task(self.abs('traces_raw.h5'),
self.abs('traces_proc.h5'),
**self.process_traces))
# link final traces to traces.h5
self.add(self._finalize_forward)
async def compute_direction(self, ws: Directory, i: int):
nprocs = getsize()
await ws.mpiexec(f'../adios/xsteepDescent kernels.bp direction.bp',
nprocs=nprocs,
walltime='adios')