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