gel.gel

  1"""Interface to forward solving functionality and file management."""
  2from .header import *
  3
  4from inspect import signature
  5from .fix_dofs_overloaded import fix_dofs
  6
  7from .geometry import *
  8from .kinematics import *
  9from .mechanics import *
 10from .helper import *
 11from .objective import *
 12
 13
 14_cell_name_cache = dict()
 15def read_cell_name(cell_data_dir):
 16    """Returns str name from directory str `cell_data_dir` with caching.
 17
 18    Will look for a file "name.txt" under provided directory and read
 19    the first line.
 20    """
 21    ret = None # Init
 22
 23    if rank == 0:
 24        # Determine full path in case of moving current directory
 25        fullpath = os.path.realpath(cell_data_dir)
 26        if fullpath in _cell_name_cache:
 27            ret = _cell_name_cache[fullpath]
 28        else:
 29            # We must read the file.
 30            with open(os.path.join(cell_data_dir, "name.txt"), "r") as fd:
 31                name = fd.readline().strip()
 32            _cell_name_cache[fullpath] = name
 33            ret = name
 34
 35    ret = comm.bcast(ret, root=0)
 36
 37    return ret
 38
 39
 40_geo_type_cache = dict()
 41def read_geo_type(cell_data_dir):
 42    """Returns str geometry type from directory str `cell_data_dir` w/ caching.
 43
 44    Will look for file "geo_type.txt" under provided directory and read
 45    the first line.
 46    """
 47    ret = None # Init
 48
 49    if rank == 0:
 50        # Determine full path in case of moving current directory
 51        fullpath = os.path.realpath(cell_data_dir)
 52        if fullpath in _geo_type_cache:
 53            ret = _geo_type_cache[fullpath]
 54        else:
 55            # We must read the file.
 56            with open(os.path.join(cell_data_dir, "geo_type.txt"), "r") as fd:
 57                geo_type = fd.readline().strip()
 58            _geo_type_cache[fullpath] = geo_type
 59            ret = geo_type
 60
 61    ret = comm.bcast(ret, root=0)
 62    
 63    return ret
 64
 65
 66def apply_fixes(
 67        geo,
 68        ctl,
 69        fillval=0.0
 70    ):
 71    """Fixes DoFs beyond the event horizon to a fixed value.
 72
 73    Uses the subdomain data in `geo` (an object of type
 74    `gel.geometry.Geometry`) to set DoFs of FEniCS function `ctl` to a
 75    uniform value of `fillval` (default 0.0) in a new function.
 76
 77    Returns a new, fixed version of `ctl`, the modulus representation,
 78    with adjoint tracing.
 79    """
 80    # Construct fixed indices first
 81    free_indices = set()
 82    dof_map = geo.V0.dofmap()
 83    my_first, my_last = dof_map.ownership_range()
 84    num_dof_this_proc = my_last - my_first
 85
 86    # Iterate through cells in the mesh
 87    for region_ind, (ci, cell) in zip(
 88        geo.dx.subdomain_data().array(),
 89        enumerate(geo.mesh.cells())
 90    ):
 91        # See if the element is in a detectable region
 92        if region_ind == geo.DETECTABLE_U:
 93            # Mark each point on that cell as free
 94            for dof in dof_map.cell_dofs(ci):
 95                free_indices.add(int(dof))
 96
 97    # Remainder should be fixed
 98    fix_indices = set(range(num_dof_this_proc)) - free_indices
 99    fix_indices = np.array(list(fix_indices))
100
101    fixed_mod_repr = fix_dofs(
102        ctl,
103        fix_indices,
104        fillval*np.ones(fix_indices.shape, dtype=float)
105    )
106
107    return fixed_mod_repr
108
109
110def get_variational_stuff(kinematics, Pi):
111    r"""Returns the right-, left-hand sides of the variational problem.
112
113    It is defined in here, as opposed to `gel.mechanics`, because the
114    operators needed to solve a nonlinear problem must be invented
115    according to the solve technique; they are not intrinsic to the
116    problem. We use a consistent-tangent Newton-Raphson solver, so we
117    need the linearized tangent operator and a right-hand-side.
118
119    * `kinematics`: object of type `gel.kinematics.Kinematics` through
120    which to differentiate with respect to displacements
121    * `Pi`: expression for energy, likely an output from `gel.mechanics`
122
123    Returns:
124    * First derivative $\frac{d\Pi}{d\mathbf{u}}$ (RHS)
125    * Second derivative $\frac{d^2\Pi}{d\mathbf{u}^2}$ (LHS)
126    """
127    geo = kinematics.geo
128
129    du = TrialFunction(geo.V)            # Incremental displacement
130    v  = TestFunction(geo.V)             # Test function
131
132    # Compute first variation of Pi 
133    # (directional derivative about u in the direction of w)
134    dPi = derivative(Pi, kinematics.u, v)
135    ddPi = derivative(dPi, kinematics.u, du)
136
137    return dPi, ddPi
138
139
140# There is no reason to record what happends here for inverse model
141@pa.tape.no_annotations
142def output_results_paraview(
143        output_folder,
144        kinematics,
145        mod_repr,
146        bx_inds=None,
147        name="simulation_output"
148    ):
149    """Outputs simulation state information as files in multiple forms.
150
151    * `output_folder`: str path to directory where files will be created
152    * `kinematics`: `gel.kinematics.Kinematics` with displacements
153    * `mod_repr`: FEniCS function with control variable DoFs
154    * `bx_inds`: dict, if not None it enables writing boundary tags on
155    hydrogel mesh with keys the tags and items the names
156    * `name`: str the name of the nodal output file
157
158    Output files:
159    * {name}.xdmf, .h5: nodal form with displacement, control variable,
160    rank ownership for DoFs, and optionally boundary conditions
161    * u_full_shape.xdmf, .h5: full shape/write_checkpoint form of
162    displacements for easy reloading back into FEniCS
163    * mod_repr_full_shape.xdmf, .h5: full shape/write_checkpoint form of
164    the control variables for easy reloading back into FEniCS
165    """
166    geo = kinematics.geo
167
168    # Helper
169    rp = ghrp(output_folder)
170
171    # Create output file for homogeneous forward model run
172    simulation_output_file = XDMFFile(rp(f"{name}.xdmf"))
173    simulation_output_file.parameters["flush_output"] = True
174    simulation_output_file.parameters["functions_share_mesh"] = True
175
176    # Renaming variables
177    kinematics.u.rename("u","displacement")
178    mod_repr.rename("mod_repr","representation of modulus")
179
180    # Rank indication
181    rank_ind = Function(geo.V0)
182    size = rank_ind.vector().local_size()
183    rank_ind.vector().set_local(rank*np.ones(size))
184    rank_ind.vector().apply("")
185    rank_ind.rename("rank","process ownership")
186
187    # Writes out the variables
188    simulation_output_file.write(kinematics.u,0)
189    simulation_output_file.write(mod_repr,0)
190    simulation_output_file.write(rank_ind,0)
191
192    # Boundary indication
193    if bx_inds:
194        for name, bx_ind in bx_inds.items():
195            # Evaluate expression, get closest fit we can (L2) with DoFs
196            indicator = project(bx_ind, geo.V0)
197            # Output result
198            indicator.rename(f"boundary_{name}", "High value on {name}")
199            simulation_output_file.write(indicator, 0)
200
201    # Full shape for creating synthetic exact target
202    shape_fname = rp("u_full_shape.xdmf")
203    save_shape(kinematics.u, shape_fname, "u")
204
205    # Full shape for multiple regularization stages
206    shape_fname = rp("mod_repr_full_shape.xdmf")
207    save_shape(mod_repr, shape_fname, "mod_repr")
208
209
210class ForwardSimulation:
211    """Minimal-prerequisite interface to running forward simulations.
212
213    Typical usage, an example of obtaining displacements from
214    homogeneous modulus:
215    ```
216    sim = ForwardSimulation("real_cell_gel", data_directory="cell_data")
217    sim.run_forward()
218    kinematics = sim.kinematics # Contains displacements u
219    ```
220
221    Many internal variables (below) are exposed for customization before
222    calls to sim.run_forward(). Be careful to keep pointers between the
223    `geo`, `kinematics`, and `mechanics` objects in sync by not
224    overwriting them or their depedency injections.
225    """
226
227    def __init__(self,
228            geometry_type,
229            d1c1=1,
230            formulation="beta",
231            mod_repr_init="zero",
232            vprint=do_nothing_print,
233            mu_ff=MU_FF_DEFAULT,
234            load_steps=1,
235            restrict_ctl_dofs_to_veh=False,
236            pc="hypre_amg",
237            tola=1e-10,
238            tolr=1e-9,
239            traction=None,
240            formulation_kwargs=dict(),
241            **kwargs
242        ):
243        r"""Constructor for object containing forward simulation state.
244
245        * `geometry_type`: str tag of what type of geometry (only
246        "real_cell_gel" -> `gel.geometry.Geometry` is currently
247        implemented)
248        * `d1c1`: float ratio $\frac{D_1}{c_1}$
249        * `formulation`: str tag of material model formulation
250        implemented in `gel.mechanics`
251        * `mod_repr_init`: (the initial value of) the control variable for
252        modulus; options in `gel.mechanics`
253        * `vprint`: callable function to perform action of print or
254        logging
255        * `mu_ff`: float far field rheometry measurement, ie.
256        $c_1=\frac{\mu}{2}$
257        * `load_steps`: int number of load steps to use in forward solve
258        (1 implies no stepping, just use the BC)
259        * `restrict_ctl_dofs_to_veh`: bool, enables fixing DoFs outside
260        event horizon for the inverse model
261        * `pc`: str name of preconditioner, ie. an option listed in
262        `list_krylov_solver_preconditioners()`
263        * `tola`: float absolute tolerance of Newton-Raphson
264        * `tolr`: float relative tolerance of Newton-Raphson
265        * `traction`: str or None. If not None, the filename of a
266        full-shape .xdmf traction function on the hydrogel mesh. Must
267        be in 1st order Lagrange space, will ignore DoFs outside of
268        surface.
269        * `formulation_kwargs`: dict of additional kwargs to the
270        material model formulation (if applicable) implemented in
271        `gel.mechanics`
272        * `kwargs`: dict of kwargs to the the geometry, for instance
273        the arguments to `gel.geometry.Geometry`
274        """
275        # Validate, get geometry
276        if traction is not None:
277            # Update kwargs -> no Dirichlet BC allowed
278            if "suppress_cell_bc" in kwargs:
279                if kwargs["suppress_cell_bc"] == False:
280                    raise ValueError("Must suppress cell BC for traction")
281
282            kwargs["suppress_cell_bc"] = True
283        self.geo = self._geo_from_kwargs(geometry_type, kwargs)
284        """Corresponding object of type `gel.geometry.Geometry`"""
285
286        # Initialize displacements, kinematic quantities
287        self.kinematics = Kinematics(self.geo)
288        """Object of type `gel.kinematics.Kinematics` with displacement
289        information
290        """
291
292        # Prepare modulus control variable
293        self.ctl = create_mod_repr(self.geo, mod_repr_init)
294        """A FEniCS function with the control variables for optimization,
295        *cf.* `mod_repr`
296        """
297
298        self.mod_repr = None
299        r"""The representation of modulus, for instance $\beta$, of type FEniCS
300        function. DoFs have already been fixed if necessary in this object.
301        *cf.* `ctl`.
302        """
303        if restrict_ctl_dofs_to_veh:
304            self.mod_repr = apply_fixes(
305                self.geo,
306                self.ctl,
307                fillval=(
308                    0.0 if (formulation in ZERO_FIX_F) else 1.0
309                )
310            )
311        else:
312            self.mod_repr = self.ctl
313
314        # Prepare mechanics
315        self.mechanics = Mechanics(
316            self.kinematics,
317            formulation,
318            mu_ff,
319            d1c1,
320            self.mod_repr,
321            **formulation_kwargs
322        )
323        """Object of type `gel.mechanics.Mechanics` with material model
324        information. Has an internal `mechanics.kinematics` variable
325        that is the same as the ForwardSimulation's `kinematics`.
326        """
327
328        self.B = Constant((0.0, 0.0, 0.0))
329        """FEniCS body force per unit volume"""
330
331        self.T = None
332        """FEniCS traction force on boundary per area in reference"""
333        if traction is None:
334            self.T = Constant((0.0, 0.0, 0.0))
335        else:
336            self.T = load_shape(self.geo.V, traction, "T")
337
338        self.solver_parameters = {
339            "newton_solver" : {
340                "absolute_tolerance" : tola,
341                "relative_tolerance" : tolr,
342                "linear_solver" : "gmres",
343                "preconditioner" : pc,
344                "maximum_iterations" : 8,
345                "error_on_nonconvergence" : False,
346                #"krylov_solver" : {
347                    #"absolute_tolerance" : 0.1*self.tola,
348                    #"relative_tolerance" : 1e-4
349                #}
350            }
351        }
352        """dict of parameters set in `NonlinearVariationalSolver`"""
353
354        self.ffc_options = {
355            "optimize": True,
356            "eliminate_zeros": True,
357            "precompute_basis_const": True,
358            "precompute_ip_const": True
359        }
360        """dict of form compiler parameters set in
361        `NonlinearVariationalProblem`
362        """
363
364        self.load_steps = load_steps
365        """Number of load steps in Dirichlet BCs during single solve"""
366
367        self.vprint = vprint
368        """Optional print/logging functionality"""
369
370    def _geo_from_kwargs(self, geometry_type, kwargs):
371        # Parse geometry
372        if geometry_type == "real_cell_gel":
373            geo_cls = Geometry
374        else:
375            raise ValueError(f"{geometry_type} unknown")
376
377        #
378        # Extract arguments
379        #
380        sig = signature(geo_cls)
381
382        # Iterate through known acceptable arguments
383        cls_args_pos_only = [ ]
384        cls_args_kwarg = dict()
385        for key, param in sig.parameters.items():
386            # Validate
387            if param.default is param.empty:
388                # Must be required
389                if key not in kwargs:
390                    raise ValueError(
391                        f"{key} must be provided to {geometry_type}"
392                    )
393
394            # Incorporate if present
395            if key in kwargs:
396                if param.kind == param.POSITIONAL_ONLY:
397                    cls_args_pos_only.append(kwargs[key])
398                else:
399                    cls_args_kwarg[key] = kwargs[key]
400
401        # Create instance
402        return geo_cls(*cls_args_pos_only, **cls_args_kwarg)
403
404    def run_forward(self):
405        r"""Solves the forward problem with saved settings.
406
407        Side effects: 
408        * `kinematics` will store solved displacements
409        * `mechanics` will contained solved displacements through the
410        same `kinematics` instance
411        * boundary conditions in `geo` are changed if load stepping is
412        enabled
413        * Uses `vprint` for logging before solve
414        """
415        # Optionally print information
416        self.vprint(f"Using D1/C1={self.mechanics.d1c1}")
417        self.vprint(f"Using far field mu={self.mechanics.mu_ff} Pa")
418        self.vprint(f"solver_parameters = {self.solver_parameters}")
419        self.vprint(f"ffc_options = {self.ffc_options}")
420
421        # Assemble variational problem with current state
422        Pi = self.mechanics.get_energy(B=self.B, T=self.T)
423        dPi, ddPi = get_variational_stuff(self.kinematics, Pi)
424
425        # Solver loop (solves the forward problem in load_steps)
426        for i in range(self.load_steps):
427            sys.stdout.flush()  
428
429            if hasattr(self.geo, "scalar"):
430                # updates the boundary conditions surface displacements
431                self.geo.scalar = (i+1)/self.load_steps
432                self.geo.update_bcs()
433
434            # Be careful of load stepping and adjoint annotation!
435            stop_tape = False
436            if i != self.load_steps-1:
437                stop_tape = True
438
439            if stop_tape:
440                pa.pause_annotation()
441
442            # Solve variational problem
443            nvp = NonlinearVariationalProblem(
444                dPi,
445                self.kinematics.u,
446                self.geo.bcs,
447                J=ddPi,
448                form_compiler_parameters=self.ffc_options
449            )
450            nvs = NonlinearVariationalSolver(nvp)
451            nvs.parameters.update(self.solver_parameters)
452            nvs.solve()
453
454            if stop_tape:
455                pa.continue_annotation()