Skip to content

Commit

Permalink
feat: update to new pymbolic
Browse files Browse the repository at this point in the history
  • Loading branch information
alexfikl committed Jan 9, 2025
1 parent 4df76e4 commit d1db7f6
Show file tree
Hide file tree
Showing 7 changed files with 259 additions and 189 deletions.
26 changes: 13 additions & 13 deletions pytential/symbolic/elasticity.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@
from functools import cached_property

import numpy as np
from pymbolic.typing import ArithmeticExpressionT
from pymbolic.typing import ArithmeticExpression
from sumpy.kernel import (AxisSourceDerivative, AxisTargetDerivative,
BiharmonicKernel, ElasticityKernel, Kernel,
LaplaceKernel, StokesletKernel, StressletKernel,
Expand Down Expand Up @@ -95,9 +95,9 @@ class ElasticityWrapperBase(ABC):

dim: int
"""Ambient dimension of the representation."""
mu: ArithmeticExpressionT
mu: ArithmeticExpression
r"""Expression or value for the shear modulus :math:`\mu`."""
nu: ArithmeticExpressionT
nu: ArithmeticExpression
r"""Expression or value for Poisson's ratio :math:`\nu`."""

@abstractmethod
Expand Down Expand Up @@ -159,9 +159,9 @@ class ElasticityDoubleLayerWrapperBase(ABC):

dim: int
"""Ambient dimension of the representation."""
mu: ArithmeticExpressionT
mu: ArithmeticExpression
r"""Expression or value for the shear modulus :math:`\mu`."""
nu: ArithmeticExpressionT
nu: ArithmeticExpression
r"""Expression or value for Poisson's ration :math:`\nu`."""

@abstractmethod
Expand Down Expand Up @@ -227,8 +227,8 @@ def _create_int_g(knl, deriv_dirs, density, **kwargs):
@dataclass
class _ElasticityWrapperNaiveOrBiharmonic:
dim: int
mu: ArithmeticExpressionT
nu: ArithmeticExpressionT
mu: ArithmeticExpression
nu: ArithmeticExpression
base_kernel: Kernel

def __post_init__(self):
Expand Down Expand Up @@ -315,8 +315,8 @@ def __init__(self, dim, mu, nu):
@dataclass
class _ElasticityDoubleLayerWrapperNaiveOrBiharmonic:
dim: int
mu: ArithmeticExpressionT
nu: ArithmeticExpressionT
mu: ArithmeticExpression
nu: ArithmeticExpression
base_kernel: Kernel

def __post_init__(self):
Expand Down Expand Up @@ -474,8 +474,8 @@ class Method(Enum):

def make_elasticity_wrapper(
dim: int,
mu: ArithmeticExpressionT = _MU_SYM_DEFAULT,
nu: ArithmeticExpressionT = _NU_SYM_DEFAULT,
mu: ArithmeticExpression = _MU_SYM_DEFAULT,
nu: ArithmeticExpression = _NU_SYM_DEFAULT,
method: Method = Method.Naive) -> ElasticityWrapperBase:
"""Creates an appropriate :class:`ElasticityWrapperBase` object.
Expand All @@ -502,8 +502,8 @@ def make_elasticity_wrapper(

def make_elasticity_double_layer_wrapper(
dim: int,
mu: ArithmeticExpressionT = _MU_SYM_DEFAULT,
nu: ArithmeticExpressionT = _NU_SYM_DEFAULT,
mu: ArithmeticExpression = _MU_SYM_DEFAULT,
nu: ArithmeticExpression = _NU_SYM_DEFAULT,
method: Method = Method.Naive) -> ElasticityDoubleLayerWrapperBase:
"""Creates an appropriate :class:`ElasticityDoubleLayerWrapperBase` object.
Expand Down
7 changes: 6 additions & 1 deletion pytential/symbolic/mappers.py
Original file line number Diff line number Diff line change
Expand Up @@ -280,7 +280,12 @@ def map_common_subexpression(self, expr):
# {{{ FlattenMapper

class FlattenMapper(FlattenMapperBase, IdentityMapper):
pass
def map_int_g(self, expr):
densities, kernel_arguments, changed = rec_int_g_arguments(self, expr)
if not changed:
return expr

return replace(expr, densities=densities, kernel_arguments=kernel_arguments)


def flatten(expr):
Expand Down
Loading

0 comments on commit d1db7f6

Please sign in to comment.