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