Homec4science
Diffusion tamaas (ci-fix-static-checks) Browse

tamaas/examples/scipy_penalty.py
40549a775601ci-fix-static-checks

Last Change
Raw File

scipy_penalty.py
View Options

import operator
import functools
import tamaas as tm
import numpy as np
import matplotlib.pyplot as plt
import scipy.optimize
from tamaas.utils import hertz_surface
class GapPenaltyFunctional(tm.Functional):
"""Quadratic penalty on negative gap functional."""
def __init__(self, penalty):
super().__init__()
self.penalty = penalty
@staticmethod
def clip(gap):
"""Clip negative gap."""
return np.clip(gap, -np.inf, 0)
def computeF(self, gap, dual):
"""Compute penalty."""
g = self.clip(gap)
return 0.5 / self.penalty * np.vdot(g, g) / g.size
def computeGradF(self, gap, gradient):
"""Compute gradient."""
gradient[:] += (1 / self.penalty) * self.clip(gap)
class ScipyContactSolver(tm.ContactSolver):
"""Gap-based penalty solver using https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.minimize.html"""
def __init__(self, model, surface, tolerance, penalty):
super().__init__(model, surface, tolerance)
if model.type != tm.model_type.basic_2d:
raise TypeError("model type is not basic_2d")
self.N = functools.reduce(operator.mul, model.shape)
self.shape = model.boundary_shape
model['gap'] = np.zeros(list(self.shape) + [1])
model.be_engine.registerDirichlet()
self.elastic = tm.ElasticFunctionalGap(
self.model.operators['Westergaard::dirichlet'], self.surface)
self.penalty = GapPenaltyFunctional(penalty)
self.addFunctionalTerm(self.elastic)
self.addFunctionalTerm(self.penalty)
def solve(self, mean_gap):
"""Solve with mean gap constraint."""
def cost(gap):
gap = gap.reshape(self.shape)
grad = np.zeros_like(gap).reshape(self.shape)
self.functional.computeGradF(gap, grad)
return self.functional.computeF(gap, grad)
def jac(gap):
gap = gap.reshape(self.shape)
grad = np.zeros_like(gap).reshape(self.shape)
self.functional.computeGradF(gap, grad)
return np.ravel(grad)
opt = scipy.optimize.minimize(
cost,
np.full(self.N, mean_gap),
jac=jac,
constraints=scipy.optimize.LinearConstraint(
np.full((1, self.N), 1 / self.N), mean_gap, mean_gap),
method='trust-constr',
tol=self.tolerance * mean_gap)
print(opt)
# Setup solved model
self.model['gap'][:] = opt.x.reshape(self.shape)
self.model.displacement[:] = self.model['gap'] + self.surface.reshape(
self.shape)
self.model.solveDirichlet()
self.model.traction[:] -= self.model.traction.min()
def main():
"""Run a simple Hertzian contact with penalty."""
N = 64
model = tm.Model(tm.model_type.basic_2d, [1, 1], [N] * 2)
surface = hertz_surface(model.system_size, model.shape, 1)
solver = ScipyContactSolver(model, surface, 1e-10, 1)
solver.solve(0.04)
plt.imshow(model.traction)
plt.colorbar()
plt.show()
if __name__ == "__main__":
main()
c4science · Help