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            u_init=None,
241            formulation_kwargs=dict(),
242            **kwargs
243        ):
244        r"""Constructor for object containing forward simulation state.
245
246        * `geometry_type`: str tag of what type of geometry (only
247        "real_cell_gel" -> `gel.geometry.Geometry` is currently
248        implemented)
249        * `d1c1`: float ratio $\frac{D_1}{c_1}$
250        * `formulation`: str tag of material model formulation
251        implemented in `gel.mechanics`
252        * `mod_repr_init`: (the initial value of) the control variable for
253        modulus; options in `gel.mechanics`
254        * `vprint`: callable function to perform action of print or
255        logging
256        * `mu_ff`: float far field rheometry measurement, ie.
257        $c_1=\frac{\mu}{2}$
258        * `load_steps`: int number of load steps to use in forward solve
259        (1 implies no stepping, just use the BC)
260        * `restrict_ctl_dofs_to_veh`: bool, enables fixing DoFs outside
261        event horizon for the inverse model
262        * `pc`: str name of preconditioner, ie. an option listed in
263        `list_krylov_solver_preconditioners()`
264        * `tola`: float absolute tolerance of Newton-Raphson
265        * `tolr`: float relative tolerance of Newton-Raphson
266        * `traction`: str or None. If not None, the filename of a
267        full-shape .xdmf traction function on the hydrogel mesh. Must
268        be in 1st order Lagrange space, will ignore DoFs outside of
269        surface. "zero" will automatically contruct a 0 traction
270        function
271        * `u_init`: str path to full-shape .xdmf file with initial
272        displacements
273        * `formulation_kwargs`: dict of additional kwargs to the
274        material model formulation (if applicable) implemented in
275        `gel.mechanics`
276        * `kwargs`: dict of kwargs to the geometry, for instance
277        the arguments to `gel.geometry.Geometry`
278        """
279        # Validate, get geometry
280        if traction is not None:
281            # Update kwargs -> no Dirichlet BC allowed
282            if "suppress_cell_bc" in kwargs:
283                if kwargs["suppress_cell_bc"] == False:
284                    raise ValueError("Must suppress cell BC for traction")
285
286            kwargs["suppress_cell_bc"] = True
287        self.geo = self._geo_from_kwargs(geometry_type, kwargs)
288        """Corresponding object of type `gel.geometry.Geometry`"""
289
290        # Initialize displacements, kinematic quantities
291        self.kinematics = None
292        """Object of type `gel.kinematics.Kinematics` with displacement
293        information
294        """
295        if u_init is None:
296            self.kinematics = Kinematics(self.geo)
297        else:
298            self.kinematics = kinematics_from_file(self.geo, u_init)
299
300        # Prepare modulus control variable
301        self.ctl = create_mod_repr(self.geo, mod_repr_init)
302        """A FEniCS function with the control variables for optimization,
303        *cf.* `mod_repr`
304        """
305
306        self.mod_repr = None
307        r"""The representation of modulus, for instance $\beta$, of type FEniCS
308        function. DoFs have already been fixed if necessary in this object.
309        *cf.* `ctl`.
310        """
311        if restrict_ctl_dofs_to_veh:
312            self.mod_repr = apply_fixes(
313                self.geo,
314                self.ctl,
315                fillval=(
316                    0.0 if (formulation in ZERO_FIX_F) else 1.0
317                )
318            )
319        else:
320            self.mod_repr = self.ctl
321
322        # Prepare mechanics
323        self.mechanics = Mechanics(
324            self.kinematics,
325            formulation,
326            mu_ff,
327            d1c1,
328            self.mod_repr,
329            **formulation_kwargs
330        )
331        """Object of type `gel.mechanics.Mechanics` with material model
332        information. Has an internal `mechanics.kinematics` variable
333        that is the same as the ForwardSimulation's `kinematics`.
334        """
335
336        self.B = Constant((0.0, 0.0, 0.0))
337        """FEniCS body force per unit volume"""
338
339        self.T = None
340        """FEniCS traction force on boundary per area in reference"""
341        if traction is None:
342            self.T = Constant((0.0, 0.0, 0.0))
343        elif traction == "zero":
344            self.T = Constant((0.0, 0.0, 0.0))
345        else:
346            self.T = load_shape(self.geo.V, traction, "T")
347
348        self.solver_parameters = {
349            "newton_solver" : {
350                "absolute_tolerance" : tola,
351                "relative_tolerance" : tolr,
352                "linear_solver" : "gmres",
353                "preconditioner" : pc,
354                "maximum_iterations" : 8,
355                "error_on_nonconvergence" : False,
356                #"krylov_solver" : {
357                    #"absolute_tolerance" : 0.1*self.tola,
358                    #"relative_tolerance" : 1e-4
359                #}
360            }
361        }
362        """dict of parameters set in `NonlinearVariationalSolver`"""
363
364        self.ffc_options = {
365            "optimize": True,
366            "eliminate_zeros": True,
367            "precompute_basis_const": True,
368            "precompute_ip_const": True
369        }
370        """dict of form compiler parameters set in
371        `NonlinearVariationalProblem`
372        """
373
374        self.load_steps = load_steps
375        """Number of load steps in Dirichlet BCs during single solve"""
376
377        self.vprint = vprint
378        """Optional print/logging functionality"""
379
380    def _geo_from_kwargs(self, geometry_type, kwargs):
381        # Parse geometry
382        if geometry_type == "real_cell_gel":
383            geo_cls = Geometry
384        else:
385            raise ValueError(f"{geometry_type} unknown")
386
387        #
388        # Extract arguments
389        #
390        sig = signature(geo_cls)
391
392        # Iterate through known acceptable arguments
393        cls_args_pos_only = [ ]
394        cls_args_kwarg = dict()
395        for key, param in sig.parameters.items():
396            # Validate
397            if param.default is param.empty:
398                # Must be required
399                if key not in kwargs:
400                    raise ValueError(
401                        f"{key} must be provided to {geometry_type}"
402                    )
403
404            # Incorporate if present
405            if key in kwargs:
406                if param.kind == param.POSITIONAL_ONLY:
407                    cls_args_pos_only.append(kwargs[key])
408                else:
409                    cls_args_kwarg[key] = kwargs[key]
410
411        # Create instance
412        return geo_cls(*cls_args_pos_only, **cls_args_kwarg)
413
414    def run_forward(self):
415        r"""Solves the forward problem with saved settings.
416
417        Side effects: 
418        * `kinematics` will store solved displacements
419        * `mechanics` will contained solved displacements through the
420        same `kinematics` instance
421        * boundary conditions in `geo` are changed if load stepping is
422        enabled
423        * Uses `vprint` for logging before solve
424        """
425        # Optionally print information
426        self.vprint(f"Using D1/C1={self.mechanics.d1c1}")
427        self.vprint(f"Using far field mu={self.mechanics.mu_ff} Pa")
428        self.vprint(f"solver_parameters = {self.solver_parameters}")
429        self.vprint(f"ffc_options = {self.ffc_options}")
430
431        # Assemble variational problem with current state
432        Pi = self.mechanics.get_energy(B=self.B, T=self.T)
433        dPi, ddPi = get_variational_stuff(self.kinematics, Pi)
434
435        # Solver loop (solves the forward problem in load_steps)
436        for i in range(self.load_steps):
437            sys.stdout.flush()  
438
439            if hasattr(self.geo, "scalar"):
440                # updates the boundary conditions surface displacements
441                self.geo.scalar = (i+1)/self.load_steps
442                self.geo.update_bcs()
443
444            # Be careful of load stepping and adjoint annotation!
445            stop_tape = False
446            if i != self.load_steps-1:
447                stop_tape = True
448
449            if stop_tape:
450                pa.pause_annotation()
451
452            # Solve variational problem
453            nvp = NonlinearVariationalProblem(
454                dPi,
455                self.kinematics.u,
456                self.geo.bcs,
457                J=ddPi,
458                form_compiler_parameters=self.ffc_options
459            )
460            nvs = NonlinearVariationalSolver(nvp)
461            nvs.parameters.update(self.solver_parameters)
462            nvs.solve()
463
464            if stop_tape:
465                pa.continue_annotation()