1"""The inverse model.
  2
  3# RECOMMENDED USAGE
  4    nohup inverse ... | tee -a out.txt &
  5
  6This will allow it to continue in the background even if the terminal
  7session dies but the user session remains. Also will log error messages
  8that may be hard-coded to stdout by FEniCS.
  9
 10# Side-Effects
 11
 12Will interact with local file system when run as a command. This
 13includes:
 14* Logs what it is doing in file specified by name `LOG_FILE`
 15* Uses a directory `results_dir` as a super-directory for all inverse
 16model results. This directory should be clear of subdirectories except
 17for those created by **specifically this program**.
 18* Creates a .csv file `inverse.csv` under `results_dir` to store a chart
 19of all options used and some result info like runtime and quantites
 20for L-curve analysis.
 21* Creates a subdirectory under `results_dir` according to the options
 22provided.
 23* In that subdirectory, stores many .xdmf files with results and
 24other information.
 25* If a subdirectory for a collection of settings already exists, will
 26not proceed forward and quietly skip.
 27
 28# API
 29"""
 30from gel import *
 31import numpy as np
 32import os
 33from sanity_utils import *
 34import dolfin as df
 35import moola
 36
 37
 38INTERM_SAVE_PERIOD = 10
 39"""Write current guess to filesytem every (this many) L-BFGS iterations"""
 40
 41
 42#############################################################################
 43### FEniCS settings
 44#############################################################################
 45parameters['linear_algebra_backend'] = 'PETSc'
 46parameters['form_compiler']['representation'] = 'uflacs'
 47parameters['form_compiler']['optimize'] = True
 48parameters['form_compiler']['cpp_optimize'] = True
 49parameters['form_compiler']['quadrature_degree'] = 3
 50
 51
 52LOG_FILE = "global_log.txt"
 53
 54
 55def _expand_obj_info(exp_info):
 56    obj_info = exp_info[5]
 57    new_exp_info = (
 58        *exp_info[:5],
 59        obj_info.gamma_num,
 60        obj_info.objective_type,
 61        obj_info.regularization_type,
 62        obj_info.objective_domain,
 63        obj_info.regularization_domain,
 64        obj_info.detectable_u_cutoff,
 65        obj_info.safe_u_weight_filename(),
 66        obj_info.apply_u_weight_to_reg,
 67        *exp_info[6:]
 68    )
 69    return new_exp_info
 70
 71
 72INVERSE_EXPERIMENT_COLS = [
 73    "Results Dir",
 74    "Cell Name",
 75    "target_file",
 76    "Formulation",
 77    "k",
 78    "gamma",
 79    "objective_type",
 80    "regularization_type",
 81    "objective_domain",
 82    "regularization_domain",
 83    "detectable_u_cutoff",
 84    "u_weight_filename",
 85    "apply_u_weight_to_reg",
 86    "lower_bound",
 87    "upper_bound",
 88    "tol",
 89    "mod_repr_init",
 90    "debug_mode",
 91    "max_iter",
 92    "load_steps",
 93    "restrict_ctl_dofs_to_veh",
 94    "preconditioner",
 95    "optimizer_backend",
 96    "bci",
 97    "bco",
 98    "outcome",
 99    "pure_obj",
100    "reg",
101    "time"
102]
103"""Names of columns in `inverse.csv`"""
104
105
106def main(args):
107    """Sets up logging to relevant files and starts the inverse model.
108
109    * `args`: `argparse.Namespace` with parsed command arguments
110
111    Side-effects: see intro to `gel.scripts.inverse`; does all the file
112    handling except the experiment's subdirectory
113
114    Computation: calls `gel_inverse`, which handles the experiment's
115    subdirectory.
116    """
117    logger = get_global_logger(LOG_FILE)
118
119    # Results
120    result_csv_path = os.path.join(args.r, "inverse.csv")
121    csv = ResultsCSV(result_csv_path, INVERSE_EXPERIMENT_COLS)
122
123    logger.info(f"Detected number of processes: {MPI.comm_world.Get_size()}")
124
125    objective_info = ObjectiveInfo(args)
126
127    exp_info = (
128        args.r,
129        args.c,
130        args.t,
131        args.f,
132        args.k,
133        objective_info,
134        *args.b,
135        args.a,
136        args.i,
137        args.debug,
138        args.m,
139        args.l,
140        not args.no_restrict_ctl_dofs_to_veh,
141        args.p,
142        args.opt_backend,
143        args.bci,
144        args.bco
145    )
146
147    logger.info(f"Beginning experiment with arguments {exp_info}.")
148
149    # Be careful of it already existing
150    try:
151        try:
152            total_time, pure_obj, reg = gel_inverse(*exp_info)
153
154            # Update DataFrame
155            exp_info = _expand_obj_info(exp_info)
156            csv.add_row(list(exp_info) + [
157                "Converged",
158                pure_obj,
159                reg,
160                float(total_time)
161            ])
162        except RuntimeError as e:
163            logger.info(
164                f"FEniCS encountered an error. Recording this outcome"
165            )
166            logger.info(e)
167
168            # Update DataFrame
169            exp_info = _expand_obj_info(exp_info)
170            csv.add_row(list(exp_info) + ["Crash", "", "", ""])
171
172        # Save into file
173        csv.save()
174    except FileExistsError:
175        logger.info(
176            "Found that experiment directory already exists. Skipping..."
177        )
178    except UnsupportedGeometryError:
179        logger.info(
180            "This geometry type is not supported by the inverse model at "
181            "this time. Skipping..."
182        )
183
184
185class UnsupportedGeometryError(Exception):
186    """Raised if `cell_data_dir` doesn't have a supported type."""
187    pass
188
189
190def gel_inverse(
191        results_dir,
192        cell_data_dir,
193        target_file,
194        formulation,
195        k,
196        objective_info,
197        lower_bound,
198        upper_bound,
199        tol,
200        mod_repr_init,
201        debug=False,
202        maxiter=10000,
203        load_steps=2,
204        restrict_ctl_dofs_to_veh=False,
205        preconditioner="hypre_amg",
206        optimizer_backend="scipy",
207        bci=None,
208        bco=None
209    ):
210    r"""Runs the inverse model, writes subdirectory and logs progress.
211
212    * `results_dir`: str path to directory to put all results
213    * `cell_data_dir`: str path to directory with geometry information
214    to use, see the constructor to `gel.geometry.Geometry` for required
215    files
216    * `target_file`: str path to full-shape .xdmf file with
217    displacements in "u"
218    * `formulation`: str a valid material model formulation, see
219    `gel.mechanics` for options
220    * `k`: float $\frac{D_1}{c_1}$ ratio
221    * `objective_info`: `gel.objective.ObjectiveInfo` with information
222    on the objective to be minimized
223    * `lower_bound`: float; either the lower bound for the "beta_tilde"
224    formulation described in `gel.mechanics` or the lower bound for
225    L-BFGS-B if the "scipy" backend if selected
226    * `upper_bound`: float; either the upper bound for the "beta_tilde"
227    formulation described in `gel.mechanics` or the upper bound for
228    L-BFGS-B if the "scipy" backend if selected
229    * `tol`: float tolerance for inverse model convergence
230    * `mod_repr_init`: str a valid modulus representation for the
231    initial guess, see `gel.mechanics` for options
232    * `debug`: bool, if True runs a Taylor test and outputs a tree of
233    the traced components of the simulation in file "graph.dot" in the
234    experiment's subdirectory
235    * `maxiter`: int maximum number of L-BFGS iterations
236    * `restrict_ctl_dofs_to_veh`: bool, it True will only control DoFs
237    in the volume within the event horizon with cutoff specified in
238    `objective_info`
239    * `preconditioner`: str preconditioner to use, see
240    `gel.gel.ForwardSimulation` for more details
241    * `optimizer_backend`: str moola for L-BFGS and that library,
242    otherwise scipy for L-BFGS-B and that library
243    * `bci`: str path to .vtk file with inner BC info, see
244    `gel.geometry.Geometry` for details
245    * `bco`: str path to .vtk file with outer BC info, see
246    `gel.geometry.Geometry` for details
247
248    Side-effects: writes many files in new subdirectory to
249    `results_dir`, see intro to `gel.scripts.inverse`
250
251    Returns:
252    * float time it took to run
253    * float objective functional value $O$ at end (see `gel.objective`)
254    * float regularization functional value $R$ at end (see
255    `gel.objective`)
256    """
257    # Input validation
258    validate_formulation(formulation)
259    validate_mod_repr_field(mod_repr_init)
260
261    # Get cell name
262    cell_name = read_cell_name(cell_data_dir)
263    # Get geometry type
264    geo_type = read_geo_type(cell_data_dir)
265    if geo_type != "real_cell_gel":
266        raise UnsupportedGeometryError(f"{geo_type} currently not supported")
267
268    # Output directory handling
269    output_dir = prepare_experiment_dir(
270        results_dir,
271        cell_name,
272        formulation,
273        restrict_ctl_dofs_to_veh,
274        mod_repr_init,
275        k,
276        objective_info.gamma_num,
277        objective_info.objective_type,
278        objective_info.regularization_type,
279        objective_info.objective_domain,
280        objective_info.regularization_domain,
281        objective_info.detectable_u_cutoff,
282        objective_info.apply_u_weight_to_reg,
283        objective_info.safe_u_weight_filename(),
284        lower_bound,
285        upper_bound,
286        tol,
287        preconditioner,
288        bci,
289        bco
290    )
291
292    # Setup logging
293    logger, logger_destructor = create_experiment_logger(output_dir)
294
295    #
296    # Initial forward run
297    #
298    logger.info(f"Beginning inverse model in {formulation} formulation.")
299    logger.info(f"Using cell data from {cell_data_dir}, '{cell_name}'")
300    logger.info(f"Using mod_repr init strat {mod_repr_init}")
301    logger.info(
302        f"Restricting mod_ctl to VEH? {restrict_ctl_dofs_to_veh}"
303    )
304    logger.info(f"Uses geometry type {geo_type}")
305    logger.info(f"Target file: {target_file}")
306    logger.info(f"Using D1/C1: {k}")
307    logger.info(f"Bounds: [{lower_bound}, {upper_bound}]")
308    logger.info(f"Using tolerance: {tol}")
309    logger.info(f"Using maximum iterations: {maxiter}")
310    logger.info(f"Using load steps: {load_steps}")
311    logger.info(f"Optimizer backend: {optimizer_backend}")
312    logger.info(f"Preconditioner: {preconditioner}")
313    logger.info(f"Override inner BC: {bci}")
314    logger.info(f"Override outer BC: {bco}")
315    objective_info.log_info(logger)
316
317    # RESET TAPE
318    set_working_tape(Tape())
319
320    # Timer start
321    timer = ExperimentTimer()
322    timer.start()
323
324    logger.info(f"Running forward model with initial guess...")
325    
326    # Deal with telling geo about detectable regions, BCs
327    addn_args = {"bci_data" : bci, "bco_data" : bco}
328    if objective_info.detectable_u_cutoff is not None:
329        addn_args["u_magnitude_subdomains_file"] = target_file
330        addn_args["detectable_u_cutoff"] = objective_info.detectable_u_cutoff
331
332    formulation_kwargs = dict()
333    if formulation == "beta_tilde":
334        formulation_kwargs["beta_min"] = lower_bound
335        formulation_kwargs["beta_max"] = upper_bound
336
337    sim = ForwardSimulation(
338        geo_type,
339        k,
340        formulation=formulation,
341        mod_repr_init=mod_repr_init,
342        load_steps=load_steps,
343        vprint=logger.info,
344        data_directory=cell_data_dir,
345        restrict_ctl_dofs_to_veh=restrict_ctl_dofs_to_veh,
346        pc=preconditioner,
347        formulation_kwargs=formulation_kwargs,
348        **addn_args
349    )
350
351    sim.run_forward()
352    kinematics, mod_repr = sim.kinematics, sim.mod_repr # Pointers
353
354    #####################
355    # Must deal with intercepted block
356    tape = get_working_tape()
357    solve_block = tape.get_blocks()[-1]
358
359    for block_var in solve_block.get_dependencies():
360        if block_var.output == kinematics.u:
361            block_var.save_output()
362            solve_block.block_var_to_save = block_var
363            solve_block.prev_soln = Function(sim.geo.V)
364    #####################
365
366    # Save
367    geo = sim.geo
368    logger.info(
369        f"Finished initial guess forward run. Saving results to {output_dir}"
370    )
371    kinematics_sim = Kinematics(geo, kinematics.u)
372    output_results_paraview(
373        output_dir,
374        kinematics_sim,
375        mod_repr,
376        name="init_guess"
377    )
378
379    # Read in target displacement
380    logger.info(f"Reading target displacements from {target_file}")
381    kinematics_target = kinematics_from_file(geo, target_file)
382
383    # Debug
384    if debug and rank == 0:
385        tape = get_working_tape()
386        tape.visualise(os.path.join(output_dir, "graph.dot"))
387    
388    #
389    # Objective function construction
390    #
391    logger.info(f"Assembling inverse objective.")
392
393    pre_assembly, pure_obj_form, reg_form = objective_info.get_objective_forms(
394        geo,
395        kinematics_target,
396        kinematics_sim,
397        mod_repr,
398        logger
399    )
400
401    # Finalize
402    obj = assemble(pre_assembly)
403
404    #
405    # Inverse model
406    #
407    # Create callback
408    i = 0
409    def callback(this_mod_repr):
410        nonlocal i, geo, kinematics, output_dir, mod_repr, solve_block
411        nonlocal optimizer_backend
412
413        if i % INTERM_SAVE_PERIOD == 0:
414            logger.info(f"Saving results for iteration {i}")
415
416            df.Function.assign(
417                kinematics.u,
418                solve_block.block_var_to_save.output
419            )
420            kinematics_mid = Kinematics(geo, kinematics.u)
421
422            # Assign values to mod_repr
423            if optimizer_backend == "scipy":
424                dm = mod_repr.function_space().dofmap()
425                local_range = dm.ownership_range()
426                mod_repr.vector().set_local(
427                    this_mod_repr[local_range[0]:local_range[1]]
428                )
429                mod_repr.vector().apply("")
430            elif optimizer_backend == "moola":
431                mod_repr.assign(this_mod_repr.data)
432
433            # Save
434            output_results_paraview(
435                output_dir,
436                kinematics_mid,
437                mod_repr,
438                name="most_recent"
439            )
440
441        i += 1
442
443    # Adjoint stuff, marking tape
444    control = Control(sim.ctl)
445    obj_hat = ReducedFunctional(obj, control)
446
447    # Check the derivative
448    if debug:
449        debug_deriv(
450            obj_hat,
451            geo,
452            mod_repr,
453            logger
454        )
455
456    # Minimize
457    logger.info(f"Calling minimization routine...")
458    #if MPI.comm_world.rank == 0:
459        #breakpoint()
460    if optimizer_backend == "scipy":
461        mod_repr_opt = minimize(
462            obj_hat,
463            method = "L-BFGS-B",
464            tol=tol,
465            bounds = (lower_bound, upper_bound),
466            options = {"maxiter":maxiter,"gtol":tol,"disp": True},
467            callback = callback
468        )
469    elif optimizer_backend == "moola":
470        problem = MoolaOptimizationProblem(obj_hat)
471        mod_repr_moola = moola.DolfinPrimalVector(sim.ctl)
472        solver = moola.BFGS(
473            problem,
474            mod_repr_moola,
475            options={
476                "jtol":tol,
477                "gtol":100*tol,
478                "Hinit":"default",
479                "maxiter":maxiter,
480                "mem_lim":10,
481                "display":2
482            },
483            hooks={"after_iteration":callback}
484        )
485        sol = solver.solve()
486        mod_repr_opt = sol["control"].data
487    else:
488        raise ValueError(f"Unknown optimizer backend {optimizer_backend}")
489
490    logger.info(f"Finished minimization")
491
492    # Assignments after optimal parameters are found
493    mod_repr.assign(mod_repr_opt)
494
495    #
496    # Post-pro: run forward of optimal params and save
497    #
498    logger.info(f"Running forward model with optimal mod_repr field...")
499
500    # Run forward model with modified mod_repr
501    sim.run_forward()
502    logger.info(f"Finished. Saving...")
503
504    # Compute QoIs
505    kinematics_post = Kinematics(geo, kinematics.u)
506
507    # Save
508    output_results_paraview(
509        output_dir,
510        kinematics_post,
511        mod_repr,
512        name="solved"
513    )
514
515    #
516    # Compute L-curve relevant quantities
517    #
518    _, pure_obj_form, reg_form = objective_info.get_objective_forms(
519        geo,
520        kinematics_target,
521        kinematics_sim,
522        mod_repr,
523        None
524    )
525
526    pure_obj = assemble(pure_obj_form)
527    if reg_form is not None:
528        reg = assemble(reg_form)
529    else:
530        reg = 0.0
531
532    #
533    # End
534    #
535    # Timer end
536    total_time = timer.end()
537    logger.info(f"Inverse modeling procedure complete")
538    logger.info(f"Took {total_time} seconds to complete")
539
540    # Tear down logging
541    logger_destructor()
542
543    return total_time, pure_obj, reg
544
545
546def inverse():
547    """The function invoked by the command. Parses arguments and passes
548    to `main`.
549    """
550    parser = get_common_parser(
551        description="One stage of the inverse model to determine modulus from "
552        "displacements"
553    )
554
555    add_objective_arguments(parser)
556
557    parser.add_argument(
558        "-r",
559        type=str,
560        metavar="RESULTS_DIR",
561        default="results",
562        help="superdirectory in which all results for given target u are stored"
563    )
564    parser.add_argument(
565        "-t",
566        type=str,
567        metavar="TARGET_U_XDMF",
568        default=None,
569        help="filename with full-shape representation of target displacements"
570    )
571    parser.add_argument(
572        "-a",
573        type=float,
574        metavar="TOL",
575        default=1e-8,
576        help="tolerance for inverse model"
577    )
578    parser.add_argument(
579        "-d",
580        "--debug",
581        action="store_true",
582        help="produces output of pyadjoint tree and performs Taylor test"
583    )
584    parser.add_argument(
585        "-m",
586        type=int,
587        metavar="MAX_ITER",
588        default=250,
589        help="maximum L-BFGS iterations"
590    )
591    parser.add_argument(
592        "-i",
593        type=str,
594        metavar="INIT_STRAT",
595        default="zero",
596        help="what to initialize the modulus representation, i.e. beta, to"
597    )
598    parser.add_argument(
599        "--no-restrict-ctl-dofs-to-veh",
600        action="store_true",
601        help="solve for the modulus everywhere in the domain instead"
602    )
603    parser.add_argument(
604        "--opt-backend",
605        type=str,
606        metavar="BACKEND",
607        choices=["scipy", "moola"],
608        default="moola",
609        help="scipy: bounded, vector repr.; moola: unbounded, L2 repr."
610    )
611
612    args = parser.parse_args()
613
614    if args.t is None:
615        args.t = os.path.join(args.c, "u_experiment.xdmf")
616
617    main(args)
618
619
620if __name__=="__main__":
621    inverse()