File size: 1,467 Bytes
4947683 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 | """Copyright (c) Meta Platforms, Inc. and affiliates."""
from __future__ import annotations
import torch
from tqdm import tqdm
from manifm.solvers import euler_step, midpoint_step, rk4_step
from flowmm.rfm.vmap import VMapManifolds
def projx_integrate_xt_to_x1(
manifold: VMapManifolds,
odefunc: callable | None,
xt: torch.Tensor,
t: torch.Tensor,
vt: torch.Tensor | None = None,
) -> torch.Tensor:
dt = 1.0 - t
if vt is None:
vt = odefunc(t, xt)
xt = euler_step(odefunc, xt, vt, t, dt)
xt = manifold.projx(xt)
return xt
@torch.no_grad()
def projx_cond_integrator_return_last(
manifold, odefunc, x0, t, method="euler", projx=True, local_coords=False, pbar=False
):
"""Has a lower memory cost since this doesn't store intermediate values."""
step_fn = {
"euler": euler_step,
# "midpoint": midpoint_step, # not supported
# "rk4": rk4_step, # not supported
}[method]
xt = x0
x1_tm1 = None # init cond
t0s = t[:-1]
if pbar:
t0s = tqdm(t0s)
for t0, t1 in zip(t0s, t[1:]):
dt = t1 - t0
vt = odefunc(t0, xt, x1_tm1)
x1_tm1 = projx_integrate_xt_to_x1(manifold, odefunc, xt, t0, vt) # new cond
x1_tm1 = manifold.abits_clamp(x1_tm1)
xt = step_fn(
odefunc, xt, vt, t0, dt, manifold=manifold if local_coords else None
)
if projx:
xt = manifold.projx(xt)
return xt
|