tum-pbs / PhiFlow

A differentiable PDE solving framework for machine learning
MIT License
1.39k stars 189 forks source link

Override boundary condition on small patch of an edge #140

Closed ljbkusters closed 10 months ago

ljbkusters commented 11 months ago

Hi there, I'm currently working on a gas-flow simulation with a partially heated pipe. I'm strugling to figure out how to set a boundary condition for a small portion of the pipe. The heat should diffuse from the heater into the gas and should then be advected by the velocity field (which itself is a viscous gas). The heater is controlled to stay at a constant temperature $T>T_0$ where $T_0$ is the temperature of the gas at the inlet.

Currently I'm working with a 2D model. I have set the velocity field extrapolation using combine_sides and have so far set the temperature extrapolation to ZERO_GRADIENT since I want to neglect temperature interactions with the unheated part of the pipe. Since no heater is added yet, the temperature of the gas remains homogenous.

# define a (staggered) verctor grid with uniform speed in x direction
velocity = flow.StaggeredGrid(
    flow.vec(x=v0, y=0),
    flow.extrapolation.combine_sides(
        x=(flow.vec(x=v0, y=0), flow.ZERO_GRADIENT), y=0),
    x=50, y=100,
    bounds=flow.Box(x=30, y=10),
    )
# define a (centered) grid with uniform temperature
temperature = flow.CenteredGrid(
    25, extrapolation=flow.ZERO_GRADIENT,
    resolution=velocity.resolution,
    bounds=velocity.bounds,
    )

Now I want to add the heater. The geometry looks something like this:

        ------------------------------------
inlet ~~>               ~>                 ~>
  v0  ~~>               ~~>                ~~>  outlet
T=T0  ~~>               ~>                 ~>
        --------------========--------------
                       Heater
                        T>T0

How do I set the boundary condition for this small patch at the bottom to a constant scalar value? I have tried looking through the phiml source coden but have not become any wiser on how to implement this. Any help would be very much apreciated. Thank you in advance.

holl- commented 11 months ago

Hi @ljbkusters, You can check out this related issue: https://github.com/tum-pbs/PhiFlow/issues/138 I'll post a custom extrapolation soon to make this easy.

Here is a workaround for 2.5.0. This sets the boundary values at the bottom center [30, 70] to 10 and 0 for the rest.

from phi.flow import *

x_pos = field.pad(CenteredGrid(0, 0, x=100, y=100), 1).points['x'].y[0]
mask = (30 < x_pos) & (x_pos < 70)
t = CenteredGrid(0, {'x': ZERO_GRADIENT, 'y': (mask * 10, 0)}, x=100, y=100)
show(field.pad(t, 1))

For future versions, I'll make it so you don't have to pad the mask beforehand.

ljbkusters commented 11 months ago

Thank you for your response. I have tried your code but it seems to throw an error once diffuse.explicit is called. Luckily this can easily be fixed, and the example then seems to be working.

import phi
print(f"USED PHIFLOW VERSION: {phi.__version__}")
print(10*"=")

from phi.flow import *

x_pos = field.pad(CenteredGrid(0, 0, x=100, y=100), 1).points['x'].y[0]
mask = (30 < x_pos) & (x_pos < 70)
t = CenteredGrid(0, {'x': ZERO_GRADIENT, 'y': (mask * 10, 0)}, x=100, y=100)
diffuse.explicit(t, 0.1, 0.1)
[expand] terminal output ``` USED PHIFLOW VERSION: 2.5.0 ========== Traceback (most recent call last): File "/root/projects/flow-simulation/heater_patch_test.py", line 10, in diffuse.explicit(t, 0.1, 0.1) File "/root/miniconda3/envs/phiflow/lib/python3.11/site-packages/phi/physics/diffuse.py", line 36, in explicit delta = laplace(field, weights=amount) if 'vector' in shape(amount) else amount * laplace(field) ^^^^^^^^^^^^^^ File "/root/miniconda3/envs/phiflow/lib/python3.11/site-packages/phi/field/_field_math.py", line 84, in laplace padded_components = [pad(field, {dim: base_widths}) for dim in axes_names] ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/root/miniconda3/envs/phiflow/lib/python3.11/site-packages/phi/field/_field_math.py", line 84, in padded_components = [pad(field, {dim: base_widths}) for dim in axes_names] ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/root/miniconda3/envs/phiflow/lib/python3.11/site-packages/phi/field/_field_math.py", line 522, in pad data = math.pad(grid.values, widths, grid.extrapolation, bounds=grid.bounds) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/root/miniconda3/envs/phiflow/lib/python3.11/site-packages/phiml/math/_ops.py", line 869, in pad result_padded = mode.pad(value, widths, **kwargs) if has_positive_widths else value ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/root/miniconda3/envs/phiflow/lib/python3.11/site-packages/phiml/math/extrapolation.py", line 1202, in pad value = ext.pad(value, ext_widths, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/root/miniconda3/envs/phiflow/lib/python3.11/site-packages/phiml/math/extrapolation.py", line 267, in pad return Extrapolation.pad(self, value, widths, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/root/miniconda3/envs/phiflow/lib/python3.11/site-packages/phiml/math/extrapolation.py", line 88, in pad values.append(self.pad_values(value, width[False], dim, False, already_padded=already_padded, **kwargs)) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/root/miniconda3/envs/phiflow/lib/python3.11/site-packages/phiml/math/extrapolation.py", line 279, in pad_values return math.expand(self.value, shape) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/root/miniconda3/envs/phiflow/lib/python3.11/site-packages/phiml/math/_magic_ops.py", line 343, in expand combined = merge_shapes(value, dims) # check that existing sizes match ^^^^^^^^^^^^^^^^^^^^^^^^^ File "/root/miniconda3/envs/phiflow/lib/python3.11/site-packages/phiml/math/_shape.py", line 1634, in merge_shapes raise IncompatibleShapes(f"Cannot merge shapes {shapes} because dimension '{dim.name}' exists with different sizes.", *shapes) phiml.math._shape.IncompatibleShapes: Cannot merge shapes [(xˢ=102), (xˢ=100, yˢ=1)] because dimension 'x' exists with different sizes. ```

I can remove the error by slicing x_pos

x_pos = field.pad(CenteredGrid(0, 0, x=100, y=100), 1).points['x'].y[0][1:-1]

This is the temperature profile after 100 steps

for _ in range(100):
    t = diffuse.explicit(t, 1., 0.1)
show(t)

image

That's neat! However this is not compatible with a parital GRADIENT_ZERO BC since that is apparently not settable with a mask. I'd be interested to see some more documentation on creating custom Extrapolation classes.

ljbkusters commented 11 months ago

I would just like to document my findings since using advect on this hacked boundary condition was giving me trouble as well. I managed to create an even more hacky version which, however, does implement a satisfactory simulation.

When using the code in the above snippet and adding an advect call I also get a shape incompatibility exception.

[expand] new code snippet ```py import phi print(f"USED PHIFLOW VERSION: {phi.__version__}") print(10*"=") from phi.flow import * x_pos = field.pad(CenteredGrid(0, 0, x=100, y=100), 1).points['x'].y[0][1:-1] mask = (30 < x_pos) & (x_pos < 70) t = CenteredGrid(0, {'x': ZERO_GRADIENT, 'y': (mask * 10, 0)}, x=100, y=100) # diffuse heat into the domain t = diffuse.explicit(t, 1., 0.1) # create a velocity field to advect the tempearture v = CenteredGrid(values=vec(x=1, y=0), extrapolation={"x": ZERO_GRADIENT, "y": (0)}, resolution=t.resolution) # advect the heat on the domain t = advect.semi_lagrangian(t, v, 1) ```
[expand] terminal output ``` USED PHIFLOW VERSION: 2.5.0 ========== Traceback (most recent call last): File "/root/projects/flow-simulation/heater_patch_test.py", line 18, in t = advect.semi_lagrangian(t, v, 1) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/root/miniconda3/envs/phiflow/lib/python3.11/site-packages/phi/physics/advect.py", line 159, in semi_lagrangian interpolated = reduce_sample(field, lookup) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/root/miniconda3/envs/phiflow/lib/python3.11/site-packages/phi/field/_field.py", line 402, in reduce_sample return field._sample(geometry, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/root/miniconda3/envs/phiflow/lib/python3.11/site-packages/phi/field/_grid.py", line 247, in _sample resampled_values = math.grid_sample(self.values, local_points, self.extrapolation, bounds=self.bounds) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/root/miniconda3/envs/phiflow/lib/python3.11/site-packages/phiml/math/_ops.py", line 947, in grid_sample result = broadcast_op(functools.partial(_grid_sample, extrap=extrap, pad_kwargs=kwargs), [grid, coordinates]) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/root/miniconda3/envs/phiflow/lib/python3.11/site-packages/phiml/math/_ops.py", line 1002, in broadcast_op return operation(*tensors) ^^^^^^^^^^^^^^^^^^^ File "/root/miniconda3/envs/phiflow/lib/python3.11/site-packages/phiml/math/_ops.py", line 967, in _grid_sample grid_padded = pad(grid, {dim: (1, 1) for dim in grid.shape.spatial.names}, extrap or e_.ZERO, **pad_kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/root/miniconda3/envs/phiflow/lib/python3.11/site-packages/phiml/math/_ops.py", line 869, in pad result_padded = mode.pad(value, widths, **kwargs) if has_positive_widths else value ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/root/miniconda3/envs/phiflow/lib/python3.11/site-packages/phiml/math/extrapolation.py", line 1202, in pad value = ext.pad(value, ext_widths, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/root/miniconda3/envs/phiflow/lib/python3.11/site-packages/phiml/math/extrapolation.py", line 267, in pad return Extrapolation.pad(self, value, widths, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/root/miniconda3/envs/phiflow/lib/python3.11/site-packages/phiml/math/extrapolation.py", line 88, in pad values.append(self.pad_values(value, width[False], dim, False, already_padded=already_padded, **kwargs)) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/root/miniconda3/envs/phiflow/lib/python3.11/site-packages/phiml/math/extrapolation.py", line 279, in pad_values return math.expand(self.value, shape) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/root/miniconda3/envs/phiflow/lib/python3.11/site-packages/phiml/math/_magic_ops.py", line 343, in expand combined = merge_shapes(value, dims) # check that existing sizes match ^^^^^^^^^^^^^^^^^^^^^^^^^ File "/root/miniconda3/envs/phiflow/lib/python3.11/site-packages/phiml/math/_shape.py", line 1634, in merge_shapes raise IncompatibleShapes(f"Cannot merge shapes {shapes} because dimension '{dim.name}' exists with different sizes.", *shapes) phiml.math._shape.IncompatibleShapes: Cannot merge shapes [(xˢ=100), (xˢ=102, yˢ=1)] because dimension 'x' exists with different sizes. ```

Removing the slicing on x_pos changes the error but does not resolve the issue

[expand] new code snippet without slicing ```py import phi print(f"USED PHIFLOW VERSION: {phi.__version__}") print(10*"=") from phi.flow import * x_pos = field.pad(CenteredGrid(0, 0, x=100, y=100), 1).points['x'].y[0] mask = (30 < x_pos) & (x_pos < 70) t = CenteredGrid(0, {'x': ZERO_GRADIENT, 'y': (mask * 10, 0)}, x=100, y=100) # diffuse heat into the domain # t = diffuse.explicit(t, 1., 0.1) # create a velocity field to advect the tempearture v = CenteredGrid(values=vec(x=1, y=0), extrapolation={"x": ZERO_GRADIENT, "y": (0)}, resolution=t.resolution) # advect the heat on the domain t = advect.semi_lagrangian(t, v, 1) ```
[expand] terminal output ``` USED PHIFLOW VERSION: 2.5.0 ========== Traceback (most recent call last): File "/root/projects/flow-simulation/heater_patch_test.py", line 18, in t = advect.semi_lagrangian(t, v, 1) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/root/miniconda3/envs/phiflow/lib/python3.11/site-packages/phi/physics/advect.py", line 159, in semi_lagrangian interpolated = reduce_sample(field, lookup) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/root/miniconda3/envs/phiflow/lib/python3.11/site-packages/phi/field/_field.py", line 402, in reduce_sample return field._sample(geometry, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/root/miniconda3/envs/phiflow/lib/python3.11/site-packages/phi/field/_grid.py", line 247, in _sample resampled_values = math.grid_sample(self.values, local_points, self.extrapolation, bounds=self.bounds) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/root/miniconda3/envs/phiflow/lib/python3.11/site-packages/phiml/math/_ops.py", line 947, in grid_sample result = broadcast_op(functools.partial(_grid_sample, extrap=extrap, pad_kwargs=kwargs), [grid, coordinates]) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/root/miniconda3/envs/phiflow/lib/python3.11/site-packages/phiml/math/_ops.py", line 1002, in broadcast_op return operation(*tensors) ^^^^^^^^^^^^^^^^^^^ File "/root/miniconda3/envs/phiflow/lib/python3.11/site-packages/phiml/math/_ops.py", line 983, in _grid_sample neighbors = _closest_grid_values(grid, coordinates, extrap or e_.ZERO, '_closest_', pad_kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/root/miniconda3/envs/phiflow/lib/python3.11/site-packages/phiml/math/_ops.py", line 905, in _closest_grid_values grid = extrap.pad(grid, non_copy_pad, **pad_kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/root/miniconda3/envs/phiflow/lib/python3.11/site-packages/phiml/math/extrapolation.py", line 1202, in pad value = ext.pad(value, ext_widths, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/root/miniconda3/envs/phiflow/lib/python3.11/site-packages/phiml/math/extrapolation.py", line 267, in pad return Extrapolation.pad(self, value, widths, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/root/miniconda3/envs/phiflow/lib/python3.11/site-packages/phiml/math/extrapolation.py", line 88, in pad values.append(self.pad_values(value, width[False], dim, False, already_padded=already_padded, **kwargs)) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/root/miniconda3/envs/phiflow/lib/python3.11/site-packages/phiml/math/extrapolation.py", line 279, in pad_values return math.expand(self.value, shape) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/root/miniconda3/envs/phiflow/lib/python3.11/site-packages/phiml/math/_magic_ops.py", line 343, in expand combined = merge_shapes(value, dims) # check that existing sizes match ^^^^^^^^^^^^^^^^^^^^^^^^^ File "/root/miniconda3/envs/phiflow/lib/python3.11/site-packages/phiml/math/_shape.py", line 1634, in merge_shapes raise IncompatibleShapes(f"Cannot merge shapes {shapes} because dimension '{dim.name}' exists with different sizes.", *shapes) phiml.math._shape.IncompatibleShapes: Cannot merge shapes [(xˢ=102), (xˢ=100, yˢ=1)] because dimension 'x' exists with different sizes. (phiflow) root@pc-luckus:~/projects/flow-simulation# ```

The only way I figured out how to circumvent this issue is by creating new grids and copying the t.values to this new grid on each diffuse and advect call where boundary conditions are set such that both functions are callable in conjunction.

import phi
print(f"USED PHIFLOW VERSION: {phi.__version__}")
print(10*"=")

from phi.flow import *

x_pos = field.pad(CenteredGrid(0, 0, x=100, y=100), 1).points['x'].y[0][1:-1]
mask = (30 < x_pos) & (x_pos < 70)

def diffusable_temperature_grid(values):
    return CenteredGrid(values,
                        {'x': ZERO_GRADIENT, 'y': (mask * 10, 0)},
                        x=100, y=100)

def advectable_temperature_grid(values):
    return CenteredGrid(values,
                        {'x': ZERO_GRADIENT, 'y': ZERO_GRADIENT},
                        x=100, y=100)

t = diffusable_temperature_grid(0)
v = CenteredGrid(values=vec(x=1, y=0),
                 extrapolation={"x": ZERO_GRADIENT, "y": (0)},
                 resolution=t.resolution)

for _ in range(100):
    t = diffuse.explicit(t, 1., 0.1)
    t = advectable_temperature_grid(t.values)
    t = advect.semi_lagrangian(t, v, 0.1)
    t = diffusable_temperature_grid(t.values)

show(t)

This produces the following heat distribution: image

Not very pretty to say the least, but working

holl- commented 11 months ago

I pushed the auto-pad of the extrapolation value to phiml/main

pip install --force-reinstall git+https://github.com/tum-pbs/PhiML.git

Now you can just write

x_pos = CenteredGrid(0, 0, x=100, y=100).points['x'].y[0]

I'll add a way of mixing extrapolations on the same boundary face shortly, e.g. ZERO_GRADIENT + const

holl- commented 11 months ago

Okay, here is an extrapolation class that can mix two extrapolations based on a mask tensor.

ConditionalExtrapolation ```python from typing import Tuple from phi.flow import * from phiml.math import extrapolation, Tensor, shape from phiml.math.extrapolation import Extrapolation class ConditionalExtrapolation(Extrapolation): def __init__(self, mask: Tensor, true_ext, false_ext): false_ext = extrapolation.as_extrapolation(false_ext) true_ext = extrapolation.as_extrapolation(true_ext) super().__init__(false_ext.pad_rank) self.mask = mask self.false_ext = false_ext self.true_ext = true_ext def to_dict(self) -> dict: raise NotImplementedError def spatial_gradient(self) -> 'Extrapolation': return ConditionalExtrapolation(self.mask, self.false_ext.spatial_gradient(), self.true_ext.spatial_gradient()) def valid_outer_faces(self, dim) -> Tuple[bool, bool]: true_lower, true_upper = self.false_ext.valid_outer_faces(dim) false_lower, false_upper = self.true_ext.valid_outer_faces(dim) return true_lower or false_lower, true_upper or false_upper @property def is_flexible(self) -> bool: return self.false_ext.is_flexible or self.true_ext.is_flexible def pad_values(self, value: Tensor, width: int, dim: str, upper_edge: bool, already_padded: dict = None, **kwargs) -> Tensor: mask = self.mask if already_padded: mask = ZERO_GRADIENT.pad(mask, already_padded) p_false = self.false_ext.pad_values(value, width, dim, upper_edge, **kwargs) p_true = self.true_ext.pad_values(value, width, dim, upper_edge, **kwargs) return math.where(mask, p_true, p_false) def pad(self, value: Tensor, widths: dict, already_padded: Optional[dict] = None, **kwargs) -> Tensor: from phiml.math._trace import ShiftLinTracer if isinstance(value, ShiftLinTracer): mask = self.mask if already_padded: mask = ZERO_GRADIENT.pad(mask, already_padded) p_false = self.false_ext.pad(value, widths, **kwargs) p_true = self.true_ext.pad(value, widths, **kwargs) return p_true * mask + p_false * (1 - mask) return Extrapolation.pad(self, value, widths, already_padded=already_padded, **kwargs) def __abs__(self): return ConditionalExtrapolation(self.mask, abs(self.false_ext), abs(self.true_ext)) def __neg__(self): return ConditionalExtrapolation(self.mask, -self.false_ext, -self.true_ext) def __add__(self, other): return ConditionalExtrapolation(self.mask, self.false_ext + other, self.true_ext + other) def __radd__(self, other): return ConditionalExtrapolation(self.mask, other + self.false_ext, other + self.true_ext) def __sub__(self, other): return ConditionalExtrapolation(self.mask, self.false_ext - other, self.true_ext - other) def __rsub__(self, other): return ConditionalExtrapolation(self.mask, other - self.false_ext, other - self.true_ext) def __mul__(self, other): return ConditionalExtrapolation(self.mask, self.false_ext * other, self.true_ext * other) def __rmul__(self, other): return ConditionalExtrapolation(self.mask, other * self.false_ext, other * self.true_ext) def __truediv__(self, other): return ConditionalExtrapolation(self.mask, self.false_ext / other, self.true_ext / other) def __rtruediv__(self, other): return ConditionalExtrapolation(self.mask, other / self.false_ext, other / self.true_ext) ```

Here is an example usage

x_pos = CenteredGrid(0, x=100, y=100).points['x'].y[0]
mask = (30 < x_pos) & (x_pos < 70)
t = CenteredGrid(Noise(), {'x': ZERO_GRADIENT, 'y': (ConditionalExtrapolation(mask, 10., ZERO_GRADIENT), ZERO_GRADIENT)}, x=100, y=100)
show(field.pad(t, 1))
ljbkusters commented 11 months ago

Thank you, I will have a look at this tomorrow.

ljbkusters commented 11 months ago

Your custom class is throwing some errors on my side. First of all I see you forgot an import of typing.Optional. Secondly I was getting a TypeError on calling field.pad(t, 1).

TypeError: __main__.ConditionalExtrapolation.pad_values() got multiple values for keyword argument 'already_padded'

I realised that the command you provided to update phiml didn't actually use the latest version you mentioned (maybe updated to the initial 1.0.1 tag?). I had to clone the phiml repo myself and install it locally.

(for anyone else copying this code, might be good to edit the comment)

On calling diffuse.explicit(t, 1., 1.) I am now getting another error

TypeError: unsupported operand type(s) for +: '_ZeroGradient' and 'ConstantExtrapolation'
Full output code ```py from typing import Tuple, Optional from phi.flow import * from phiml.math import extrapolation, Tensor, shape from phiml.math.extrapolation import Extrapolation class ConditionalExtrapolation(Extrapolation): def __init__(self, mask: Tensor, true_ext, false_ext): false_ext = extrapolation.as_extrapolation(false_ext) true_ext = extrapolation.as_extrapolation(true_ext) super().__init__(false_ext.pad_rank) self.mask = mask self.false_ext = false_ext self.true_ext = true_ext def to_dict(self) -> dict: raise NotImplementedError def spatial_gradient(self) -> 'Extrapolation': return ConditionalExtrapolation(self.mask, self.false_ext.spatial_gradient(), self.true_ext.spatial_gradient()) def valid_outer_faces(self, dim) -> Tuple[bool, bool]: true_lower, true_upper = self.false_ext.valid_outer_faces(dim) false_lower, false_upper = self.true_ext.valid_outer_faces(dim) return true_lower or false_lower, true_upper or false_upper @property def is_flexible(self) -> bool: return self.false_ext.is_flexible or self.true_ext.is_flexible def pad_values(self, value: Tensor, width: int, dim: str, upper_edge: bool, already_padded: dict = None, **kwargs) -> Tensor: mask = self.mask if already_padded: mask = ZERO_GRADIENT.pad(mask, already_padded) p_false = self.false_ext.pad_values(value, width, dim, upper_edge, **kwargs) p_true = self.true_ext.pad_values(value, width, dim, upper_edge, **kwargs) return math.where(mask, p_true, p_false) def pad(self, value: Tensor, widths: dict, already_padded: Optional[dict] = None, **kwargs) -> Tensor: from phiml.math._trace import ShiftLinTracer if isinstance(value, ShiftLinTracer): mask = self.mask if already_padded: mask = ZERO_GRADIENT.pad(mask, already_padded) p_false = self.false_ext.pad(value, widths, **kwargs) p_true = self.true_ext.pad(value, widths, **kwargs) return p_true * mask + p_false * (1 - mask) return Extrapolation.pad(self, value, widths, already_padded=already_padded, **kwargs) def __abs__(self): return ConditionalExtrapolation(self.mask, abs(self.false_ext), abs(self.true_ext)) def __neg__(self): return ConditionalExtrapolation(self.mask, -self.false_ext, -self.true_ext) def __add__(self, other): return ConditionalExtrapolation(self.mask, self.false_ext + other, self.true_ext + other) def __radd__(self, other): return ConditionalExtrapolation(self.mask, other + self.false_ext, other + self.true_ext) def __sub__(self, other): return ConditionalExtrapolation(self.mask, self.false_ext - other, self.true_ext - other) def __rsub__(self, other): return ConditionalExtrapolation(self.mask, other - self.false_ext, other - self.true_ext) def __mul__(self, other): return ConditionalExtrapolation(self.mask, self.false_ext * other, self.true_ext * other) def __rmul__(self, other): return ConditionalExtrapolation(self.mask, other * self.false_ext, other * self.true_ext) def __truediv__(self, other): return ConditionalExtrapolation(self.mask, self.false_ext / other, self.true_ext / other) def __rtruediv__(self, other): return ConditionalExtrapolation(self.mask, other / self.false_ext, other / self.true_ext) if __name__ == "__main__": x_pos = CenteredGrid(0, x=100, y=100).points['x'].y[0] mask = (30 < x_pos) & (x_pos < 70) t = CenteredGrid(Noise(), {'x': ZERO_GRADIENT, 'y': (ConditionalExtrapolation(mask, 10., ZERO_GRADIENT), ZERO_GRADIENT)}, x=100, y=100) show(field.pad(t, 1)) diffuse.explicit(t, 1., 1.) ``` terminal output ``` Traceback (most recent call last): File "/root/projects/flow-simulation/custom_extrapolation_heater_patch.py", line 87, in diffuse.explicit(t, 1., 1.) File "/root/miniconda3/envs/phiflow/lib/python3.11/site-packages/phi/physics/diffuse.py", line 36, in explicit delta = laplace(field, weights=amount) if 'vector' in shape(amount) else amount * laplace(field) ^^^^^^^^^^^^^^ File "/root/miniconda3/envs/phiflow/lib/python3.11/site-packages/phi/field/_field_math.py", line 86, in laplace result_components = [sum([value * shift_ for value, shift_ in zip(values, shifted_component)]) / field.dx.vector[dim]**2 for shifted_component, dim in zip(shifted_components, axes_names)] ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/root/miniconda3/envs/phiflow/lib/python3.11/site-packages/phi/field/_field_math.py", line 86, in result_components = [sum([value * shift_ for value, shift_ in zip(values, shifted_component)]) / field.dx.vector[dim]**2 for shifted_component, dim in zip(shifted_components, axes_names)] ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/root/miniconda3/envs/phiflow/lib/python3.11/site-packages/phi/field/_field.py", line 262, in __add__ return self._op2(other, lambda d1, d2: d1 + d2) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/root/miniconda3/envs/phiflow/lib/python3.11/site-packages/phi/field/_field.py", line 307, in _op2 extrapolation_ = operator(self._extrapolation, other.extrapolation) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/root/miniconda3/envs/phiflow/lib/python3.11/site-packages/phi/field/_field.py", line 262, in return self._op2(other, lambda d1, d2: d1 + d2) ~~~^~~~ File "/root/miniconda3/envs/phiflow/lib/python3.11/site-packages/phiml/math/extrapolation.py", line 1250, in __add__ return self._op2(other, lambda e1, e2: e1 + e2) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/root/miniconda3/envs/phiflow/lib/python3.11/site-packages/phiml/math/extrapolation.py", line 1276, in _op2 return combine_sides({k: operator(ext, other.ext[k]) for k, ext in self.ext.items()}) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/root/miniconda3/envs/phiflow/lib/python3.11/site-packages/phiml/math/extrapolation.py", line 1276, in return combine_sides({k: operator(ext, other.ext[k]) for k, ext in self.ext.items()}) ^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/root/miniconda3/envs/phiflow/lib/python3.11/site-packages/phiml/math/extrapolation.py", line 1250, in return self._op2(other, lambda e1, e2: e1 + e2) ~~~^~~~ File "/root/projects/flow-simulation/custom_extrapolation_heater_patch.py", line 58, in __add__ return ConditionalExtrapolation(self.mask, self.false_ext + other, self.true_ext + other) ~~~~~~~~~~~~~~~^~~~~~~ File "/root/projects/flow-simulation/custom_extrapolation_heater_patch.py", line 61, in __radd__ return ConditionalExtrapolation(self.mask, other + self.false_ext, other + self.true_ext) ~~~~~~^~~~~~~~~~~~~~~ TypeError: unsupported operand type(s) for +: '_ZeroGradient' and 'ConstantExtrapolation' ```

Are you getting the same error?

holl- commented 11 months ago

Yes sorry, here is a fix for the extrapolation class:

ConditionalExtrapolation ```python from typing import Tuple, Optional, Callable from phi.flow import * from phiml.math import extrapolation, Tensor, shape from phiml.math.extrapolation import Extrapolation class ConditionalExtrapolation(Extrapolation): def __init__(self, mask: Tensor, true_ext, false_ext): true_ext = extrapolation.as_extrapolation(true_ext) false_ext = extrapolation.as_extrapolation(false_ext) super().__init__(false_ext.pad_rank) self.mask = mask self.false_ext = false_ext self.true_ext = true_ext def to_dict(self) -> dict: raise NotImplementedError def spatial_gradient(self) -> 'Extrapolation': return ConditionalExtrapolation(self.mask, self.true_ext.spatial_gradient(), self.false_ext.spatial_gradient()) def valid_outer_faces(self, dim) -> Tuple[bool, bool]: true_lower, true_upper = self.false_ext.valid_outer_faces(dim) false_lower, false_upper = self.true_ext.valid_outer_faces(dim) return true_lower or false_lower, true_upper or false_upper @property def is_flexible(self) -> bool: return self.false_ext.is_flexible or self.true_ext.is_flexible def pad_values(self, value: Tensor, width: int, dim: str, upper_edge: bool, already_padded: dict = None, **kwargs) -> Tensor: mask = self.mask if already_padded: mask = ZERO_GRADIENT.pad(mask, already_padded) p_false = self.false_ext.pad_values(value, width, dim, upper_edge, **kwargs) p_true = self.true_ext.pad_values(value, width, dim, upper_edge, **kwargs) return math.where(mask, p_true, p_false) def pad(self, value: Tensor, widths: dict, already_padded: Optional[dict] = None, **kwargs) -> Tensor: from phiml.math._trace import ShiftLinTracer if isinstance(value, ShiftLinTracer): mask = self.mask if already_padded: mask = ZERO_GRADIENT.pad(mask, already_padded) p_false = self.false_ext.pad(value, widths, **kwargs) p_true = self.true_ext.pad(value, widths, **kwargs) return p_true * mask + p_false * (1 - mask) return Extrapolation.pad(self, value, widths, already_padded=already_padded, **kwargs) def __abs__(self): return ConditionalExtrapolation(self.mask, abs(self.true_ext), abs(self.false_ext)) def __neg__(self): return ConditionalExtrapolation(self.mask, -self.true_ext, -self.false_ext) def _op2(self, other, op: Callable): if isinstance(other, ConditionalExtrapolation) and math.close(other.mask, self.mask): return ConditionalExtrapolation(self.mask, op(self.true_ext, other.true_ext), op(self.false_ext, other.false_ext)) return ConditionalExtrapolation(self.mask, op(self.true_ext, other), op(self.false_ext, other)) def __add__(self, other): return self._op2(other, lambda x, y: x + y) def __radd__(self, other): return self._op2(other, lambda y, x: x + y) def __sub__(self, other): return self._op2(other, lambda x, y: x - y) def __rsub__(self, other): return self._op2(other, lambda y, x: x - y) def __mul__(self, other): return self._op2(other, lambda x, y: x * y) def __rmul__(self, other): return self._op2(other, lambda y, x: x * y) def __truediv__(self, other): return self._op2(other, lambda x, y: x / y) def __rtruediv__(self, other): return self._op2(other, lambda y, x: x / y) ```

To make it work with advection, you'll need to pull the latest phiml again.