feat: data file exporter node
As with many things, there seems to be an obvious convergent design philosophy here wrt. data flow. We should definitely be using `SimSymbol` for a lot more things; it solves a lot of the pain points related to figuring out what on Earth should go into the `InfoFlow` in which situations. We desperately need to iron out the `*Flow` object semantics The surprise MVP of the day is `Polars`. What a gorgeous and fast dataframe library. We initially wrote it off as being unsuited to multidimensional data, but case in point, a whole lot of useful data can indeed be expressed as 2D. For all of these cases, be it loading/saving or processing, `Polars` is truly an ideal choice. Work continues.main
parent
0f2f494868
commit
f5d19abecd
GPG Key ID:
AD901CB0F3701434
|
|
@ -21,11 +21,12 @@ dependencies = [
|
|||
# Pin Blender 4.1.0-Compatible Versions
|
||||
## The dependency resolver will report if anything is wonky.
|
||||
"urllib3==1.26.8",
|
||||
#"requests==2.27.1", ## Conflict with dev-dep commitizen
|
||||
#"requests==2.27.1", ## Conflict with dev-dep commitizen
|
||||
"numpy==1.24.3",
|
||||
"idna==3.3",
|
||||
#"charset-normalizer==2.0.10", ## Conflict with dev-dep commitizen
|
||||
#"charset-normalizer==2.0.10", ## Conflict with dev-dep commitizen
|
||||
"certifi==2021.10.8",
|
||||
"polars>=0.20.26",
|
||||
]
|
||||
## When it comes to dev-dep conflicts:
|
||||
## -> It's okay to leave Blender-pinned deps out of prod; Blender still has them.
|
||||
|
|
|
|||
|
|
@ -123,6 +123,7 @@ pillow==10.2.0
|
|||
# via matplotlib
|
||||
platformdirs==4.2.1
|
||||
# via virtualenv
|
||||
polars==0.20.26
|
||||
pre-commit==3.7.0
|
||||
prompt-toolkit==3.0.36
|
||||
# via questionary
|
||||
|
|
|
|||
|
|
@ -96,6 +96,7 @@ partd==1.4.1
|
|||
# via dask
|
||||
pillow==10.2.0
|
||||
# via matplotlib
|
||||
polars==0.20.26
|
||||
pydantic==2.7.1
|
||||
# via tidy3d
|
||||
pydantic-core==2.18.2
|
||||
|
|
|
|||
|
|
@ -41,6 +41,9 @@ class OperatorType(enum.StrEnum):
|
|||
SocketCloudAuthenticate = enum.auto()
|
||||
SocketReloadCloudFolderList = enum.auto()
|
||||
|
||||
# Node: ExportDataFile
|
||||
NodeExportDataFile = enum.auto()
|
||||
|
||||
# Node: Tidy3DWebImporter
|
||||
NodeLoadCloudSim = enum.auto()
|
||||
|
||||
|
|
|
|||
|
|
@ -59,6 +59,7 @@ from .mobj_types import ManagedObjType
|
|||
from .node_types import NodeType
|
||||
from .sim_types import (
|
||||
BoundCondType,
|
||||
DataFileFormat,
|
||||
NewSimCloudTask,
|
||||
SimAxisDir,
|
||||
SimFieldPols,
|
||||
|
|
@ -103,6 +104,7 @@ __all__ = [
|
|||
'BLSocketType',
|
||||
'NodeType',
|
||||
'BoundCondType',
|
||||
'DataFileFormat',
|
||||
'NewSimCloudTask',
|
||||
'SimAxisDir',
|
||||
'SimFieldPols',
|
||||
|
|
|
|||
|
|
@ -31,17 +31,39 @@ LazyFunction: typ.TypeAlias = typ.Callable[[typ.Any, ...], typ.Any]
|
|||
|
||||
@dataclasses.dataclass(frozen=True, kw_only=True)
|
||||
class LazyValueFuncFlow:
|
||||
r"""Wraps a composable function, providing useful information and operations.
|
||||
r"""Defines a flow of data as incremental function composition.
|
||||
|
||||
# Data Flow as Function Composition
|
||||
When using nodes to do math, it can be a good idea to express a **flow of data as the composition of functions**.
|
||||
For specific math system usage instructions, please consult the documentation of relevant nodes.
|
||||
|
||||
Each node creates a new function, which uses the still-unknown (aka. **lazy**) output of the previous function to plan some calculations.
|
||||
Some new arguments may also be added, of course.
|
||||
# Introduction
|
||||
When using nodes to do math, it becomes immediately obvious to express **flows of data as composed function chains**.
|
||||
Doing so has several advantages:
|
||||
|
||||
## Root Function
|
||||
Of course, one needs to select a "bottom" function, which has no previous function as input.
|
||||
Thus, the first step is to define this **root function**:
|
||||
- **Interactive**: Since no large-array math is being done, the UI can be designed to feel fast and snappy.
|
||||
- **Symbolic**: Since no numerical math is being done yet, we can choose to keep our input parameters as symbolic variables with no performance impact.
|
||||
- **Performant**: Since no operations are happening, the UI feels fast and snappy.
|
||||
|
||||
## Strongly Related FlowKinds
|
||||
For doing math, `LazyValueFunc` relies on two other `FlowKind`s, which must run in parallel:
|
||||
|
||||
- `FlowKind.Info`: Tracks the name, `spux.MathType`, unit (if any), length, and index coordinates for the raw data object produced by `LazyValueFunc`.
|
||||
- `FlowKind.Params`: Tracks the particular values of input parameters to the lazy function, each of which can also be symbolic.
|
||||
|
||||
For more, please see the documentation for each.
|
||||
|
||||
## Non-Mathematical Use
|
||||
Of course, there are many interesting uses of incremental function composition that aren't mathematical.
|
||||
|
||||
For such cases, the usage is identical, but the complexity is lessened; for example, `Info` no longer effectively needs to flow in parallel.
|
||||
|
||||
|
||||
|
||||
# Lazy Math: Theoretical Foundation
|
||||
This `FlowKind` is the critical component of a functional-inspired system for lazy multilinear math.
|
||||
Thus, it makes sense to describe the math system here.
|
||||
|
||||
## `depth=0`: Root Function
|
||||
To start a composition chain, a function with no inputs must be defined as the "root", or "bottom".
|
||||
|
||||
$$
|
||||
f_0:\ \ \ \ \biggl(
|
||||
|
|
@ -55,7 +77,7 @@ class LazyValueFuncFlow:
|
|||
\biggr) \to \text{output}_0
|
||||
$$
|
||||
|
||||
We'll express this simple snippet like so:
|
||||
In Python, such a construction would look like this:
|
||||
|
||||
```python
|
||||
# Presume 'A0', 'KV0' contain only the args/kwargs for f_0
|
||||
|
|
@ -70,11 +92,9 @@ class LazyValueFuncFlow:
|
|||
output_0 = lazy_value_func.func(*A0_computed, **KV0_computed)
|
||||
```
|
||||
|
||||
So far so good.
|
||||
But of course, nothing interesting has really happened yet.
|
||||
|
||||
## Composing Functions
|
||||
The key thing is the next step: The function that uses the result of $f_0$!
|
||||
## `depth>0`: Composition Chaining
|
||||
So far, so easy.
|
||||
Now, let's add a function that uses the result of $f_0$, without yet computing it.
|
||||
|
||||
$$
|
||||
f_1:\ \ \ \ \biggl(
|
||||
|
|
@ -86,10 +106,14 @@ class LazyValueFuncFlow:
|
|||
\biggr) \to \text{output}_1
|
||||
$$
|
||||
|
||||
Notice that _$f_1$ needs the arguments of both $f_0$ and $f_1$_.
|
||||
Tracking arguments is already getting out of hand; we already have to use `...` to keep it readeable!
|
||||
Note:
|
||||
- $f_1$ must take the arguments of both $f_0$ and $f_1$.
|
||||
- The complexity is getting notationally complex; we already have to use `...` to represent "the last function's arguments".
|
||||
|
||||
But doing so with `LazyValueFunc` is not so complex:
|
||||
In other words, **there's suddenly a lot to manage**.
|
||||
Even worse, the bigger the $n$, the more complexity we must real with.
|
||||
|
||||
This is where the Python version starts to show its purpose:
|
||||
|
||||
```python
|
||||
# Presume 'A1', 'K1' contain only the args/kwarg names for f_1
|
||||
|
|
@ -107,46 +131,114 @@ class LazyValueFuncFlow:
|
|||
output_1 = lazy_value_func_1.func(*A_computed, **KW_computed)
|
||||
```
|
||||
|
||||
We only need the arguments to $f_1$, and `LazyValueFunc` figures out how to make one function with enough arguments to call both.
|
||||
By using `LazyValueFunc`, we've guaranteed that even hugely deep $n$s won't ever look more complicated than this.
|
||||
|
||||
## Isn't Laying Functions Slow/Hard?
|
||||
Imagine that each function represents the action of a node, each of which performs expensive calculations on huge `numpy` arrays (**as one does when processing electromagnetic field data**).
|
||||
At the end, a node might run the entire procedure with all arguments:
|
||||
## `max depth`: "Realization"
|
||||
So, we've composed a bunch of functions of functions of ...
|
||||
We've also tracked their arguments, either manually (as above), or with the help of a handy `ParamsFlow` object.
|
||||
|
||||
But it'd be pointless to just compose away forever.
|
||||
We do actually need the data that they claim to compute now:
|
||||
|
||||
```python
|
||||
# A_all and KW_all must be tracked on the side.
|
||||
output_n = lazy_value_func_n.func(*A_all, **KW_all)
|
||||
```
|
||||
|
||||
It's rough: Most non-trivial pipelines drown in the time/memory overhead of incremental `numpy` operations - individually fast, but collectively iffy.
|
||||
Of course, this comes with enormous overhead.
|
||||
Aside from the function calls themselves (which can be non-trivial), we must also contend with the enormous inefficiency of performing array operations sequentially.
|
||||
|
||||
The killer feature of `LazyValueFuncFlow` is a sprinkle of black magic:
|
||||
That brings us to the killer feature of `LazyValueFuncFlow`, and the motivating reason for doing any of this at all:
|
||||
|
||||
```python
|
||||
func_n_jax = lazy_value_func_n.func_jax
|
||||
output_n = func_n_jax(*A_all, **KW_all) ## Runs on your GPU
|
||||
output_n = lazy_value_func_n.func_jax(*A_all, **KW_all)
|
||||
```
|
||||
|
||||
What happened was, **the entire pipeline** was compiled and optimized for high performance on not just your CPU, _but also (possibly) your GPU_.
|
||||
All the layered function calls and inefficient incremental processing is **transformed into a high-performance program**.
|
||||
What happened was, **the entire pipeline** was compiled, optimized, and computed with bare-metal performance on either a CPU, GPU, or TPU.
|
||||
With the help of the `jax` library (and its underlying OpenXLA bytecode), all of that inefficiency has been optimized based on _what we're trying to do_, not _exactly how we're doing it_, in order to maximize the use of modern massively-parallel devices.
|
||||
|
||||
Thank `jax` - specifically, `jax.jit` (https://jax.readthedocs.io/en/latest/_autosummary/jax.jit.html#jax.jit), which internally enables this magic with a single function call.
|
||||
See the documentation of `LazyValueFunc.func_jax()` for more information on this process.
|
||||
|
||||
## Other Considerations
|
||||
**Auto-Differentiation**: Incredibly, `jax.jit` isn't the killer feature of `jax`. The function that comes out of `LazyValueFuncFlow` can also be differentiated with `jax.grad` (read: high-performance Jacobians for optimizing input parameters).
|
||||
|
||||
Though designed for machine learning, there's no reason other fields can't enjoy their inventions!
|
||||
|
||||
**Impact of Independent Caching**: JIT'ing can be slow.
|
||||
That's why `LazyValueFuncFlow` has its own `FlowKind` "lane", which means that **only changes to the processing procedures will cause recompilation**.
|
||||
# Lazy Math: Practical Considerations
|
||||
By using nodes to express a lazily-composed chain of mathematical operations on tensor-like data, we strike a difficult balance between UX, flexibility, and performance.
|
||||
|
||||
Generally, adjustable values that affect the output will flow via the `Param` "lane", which has its own incremental caching, and only meets the compiled function when it's "plugged in" for final evaluation.
|
||||
The effect is a feeling of snappiness and interactivity, even as the volume of data grows.
|
||||
## UX
|
||||
UX is often more a matter of art/taste than science, so don't trust these philosophies too much - a lot of the analysis is entirely personal and subjective.
|
||||
|
||||
The goal for our UX is to minimize the "frictions" that cause cascading, small-scale _user anxiety_.
|
||||
|
||||
Especially of concern in a visual math system on large data volumes is **UX latency** - also known as **lag**.
|
||||
In particular, the most important facet to minimize is _emotional burden_ rather than quantitative milliseconds.
|
||||
Any repeated moment-to-moment friction can be very damaging to a user's ability to be productive in a piece of software.
|
||||
|
||||
Unfortunately, in a node-based architecture, data must generally be computed one step at a time, whenever any part of it is needed, and it must do so before any feedback can be provided.
|
||||
In a math system like this, that data is presumed "big", and as such we're left with the unfortunate experience of even the most well-cached, high-performance operations causing _just about anything_ to **feel** like a highly unpleasant slog as soon as the data gets big enough.
|
||||
**This discomfort scales with the size of data**, by the way, which might just cause users to never even attempt working with the data volume that they actually need.
|
||||
|
||||
For electrodynamic field analysis, it's not uncommon for toy examples to expend hundreds of megabytes of memory, all of which needs all manner of interesting things done to it.
|
||||
It can therefore be very easy to stumble across that feeling of "slogging through" any program that does real-world EM field analysis.
|
||||
This has consequences: The user tries fewer ideas, becomes more easily frustrated, and might ultimately accomplish less.
|
||||
|
||||
Lazy evaluation allows _delaying_ a computation to a point in time where the user both expects and understands the time that the computation takes.
|
||||
For example, the user experience of pressing a button clearly marked with terminology like "load", "save", "compute", "run", seems to be paired to a greatly increased emotional tolerance towards the latency introduced by pressing that button (so long as it is only clickable when it works).
|
||||
To a lesser degree, attaching a node link also seems to have this property, though that tolerance seems to fall as proficiency with the node-based tool rises.
|
||||
As a more nuanced example, when lag occurs due to the computing an image-based plot based on live-computed math, then the visual feedback of _the plot actually changing_ seems to have a similar effect, not least because it's emotionally well-understood that detaching the `Viewer` node would also remove the lag.
|
||||
|
||||
In short: Even if lazy evaluation didn't make any math faster, it will still _feel_ faster (to a point - raw performance obviously still matters).
|
||||
Without `LazyValueFuncFlow`, the point of evaluation cannot be chosen at all, which is a huge issue for all the named reasons.
|
||||
With `LazyValueFuncFlow`, better-chosen evaluation points can be chosen to cause the _user experience_ of high performance, simply because we were able to shift the exact same computation to a point in time where the user either understands or tolerates the delay better.
|
||||
|
||||
## Flexibility
|
||||
Large-scale math is done on tensors, whether one knows (or likes!) it or not.
|
||||
To this end, the indexed arrays produced by `LazyValueFuncFlow.func_jax` aren't quite sufficient for most operations we want to do:
|
||||
|
||||
- **Naming**: What _is_ each axis?
|
||||
Unnamed index axes are sometimes easy to decode, but in general, names have an unexpectedly critical function when operating on arrays.
|
||||
Lack of names is a huge part of why perfectly elegant array math in ex. `MATLAB` or `numpy` can so easily feel so incredibly convoluted.
|
||||
_Sometimes arrays with named axes are called "structured arrays".
|
||||
|
||||
- **Coordinates**: What do the indices of each axis really _mean_?
|
||||
For example, an array of $500$ by-wavelength observations of power (watts) can't be limited to between $200nm$ to $700nm$.
|
||||
But they can be limited to between index `23` to `298`.
|
||||
I'm **just supposed to know** that `23` means $200nm$, and that `298` indicates the observation just after $700nm$, and _hope_ that this is exact enough.
|
||||
|
||||
Not only do we endeavor to track these, but we also introduce unit-awareness to the coordinates, and design the entire math system to visually communicate the state of arrays before/after every single computation, as well as only expose operations that this tracked data indicates possible.
|
||||
|
||||
In practice, this happens in `FlowKind.Info`, which due to having its own `FlowKind` "lane" can be adjusted without triggering changes to (and therefore recompilation of) the `FlowKind.LazyValueFunc` chain.
|
||||
**Please consult the `InfoFlow` documentation for more**.
|
||||
|
||||
## Performance
|
||||
All values introduced while processing are kept in a seperate `FlowKind` lane, with its own incremental caching: `FlowKind.Params`.
|
||||
|
||||
It's a simple mechanism, but for the cost of introducing an extra `FlowKind` "lane", all of the values used to process data can be live-adjusted without the overhead of recompiling the entire `LazyValueFunc` every time anything changes.
|
||||
Moreover, values used to process data don't even have to be numbers yet: They can be expressions of symbolic variables, complete with units, which are only realized at the very end of the chain, by the node that absolutely cannot function without the actual numerical data.
|
||||
|
||||
See the `ParamFlow` documentation for more information.
|
||||
|
||||
|
||||
|
||||
# Conclusion
|
||||
There is, of course, a lot more to say about the math system in general.
|
||||
A few teasers of what nodes can do with this system:
|
||||
|
||||
**Auto-Differentiation**: `jax.jit` isn't even really the killer feature of `jax`.
|
||||
`jax` can automatically differentiate `LazyValueFuncFlow.func_jax` with respect to any input parameter, including for fwd/bck jacobians/hessians, with robust numerical stability.
|
||||
When used in
|
||||
**Symbolic Interop**: Any `sympy` expression containing symbolic variables can be compiled, by `sympy`, into a `jax`-compatible function which takes
|
||||
We make use of this in the `Expr` socket, enabling true symbolic math to be used in high-performance lazy `jax` computations.
|
||||
**Tidy3D Interop**: For some parameters of some simulation objects, `tidy3d` actually supports adjoint-driven differentiation _through the cloud simulation_.
|
||||
This enables our humble interface to implement fully functional **inverse design** of parameterized structures, using only nodes.
|
||||
|
||||
But above all, we hope that this math system is fun, practical, and maybe even interesting.
|
||||
|
||||
Attributes:
|
||||
func: The function that the object encapsulates.
|
||||
bound_args: Arguments that will be packaged into function, which can't be later modifier.
|
||||
func_kwargs: Arguments to be specified by the user at the time of use.
|
||||
supports_jax: Whether the contained `self.function` can be compiled with JAX's JIT compiler.
|
||||
func: The function that generates the represented value.
|
||||
func_args: The constrained identity of all positional arguments to the function.
|
||||
func_kwargs: The constrained identity of all keyword arguments to the function.
|
||||
supports_jax: Whether `self.func` can be compiled with JAX's JIT compiler.
|
||||
See the documentation of `self.func_jax()`.
|
||||
"""
|
||||
|
||||
func: LazyFunction
|
||||
|
|
@ -156,13 +248,160 @@ class LazyValueFuncFlow:
|
|||
func_kwargs: dict[str, spux.MathType | spux.PhysicalType] = dataclasses.field(
|
||||
default_factory=dict
|
||||
)
|
||||
## TODO: Use SimSymbol instead of the MathType|PT union.
|
||||
## -- SimSymbol is an ideal pivot point for both, as well as valid domains.
|
||||
## -- SimSymbol has more semantic meaning, including a name.
|
||||
## -- If desired, SimSymbols could maybe even require a specific unit.
|
||||
## It could greatly simplify a whole lot of pain associated with func_args.
|
||||
supports_jax: bool = False
|
||||
|
||||
# Merging
|
||||
####################
|
||||
# - Functions
|
||||
####################
|
||||
@functools.cached_property
|
||||
def func_jax(self) -> LazyFunction:
|
||||
"""Compile `self.func` into optimized XLA bytecode using `jax.jit`.
|
||||
|
||||
Not all functions can be compiled like this by `jax`.
|
||||
A few critical criteria include:
|
||||
|
||||
- **Only JAX Ops**: All operations performed within the function must be explicitly compatible with `jax`, which generally means only using functions in `jax.lax`, `jax.numpy`
|
||||
- **Known Shape**: The exact dimensions of the output, and of the inputs, must be known at `jit`-time.
|
||||
|
||||
In return, one receives:
|
||||
|
||||
- **Automatic Differentiation**: `jax` can robustly differentiate this function with respect to _any_ parameter.
|
||||
This includes Jacobians and Hessians, forwards and backwards, real or complex, all with good numerical stability.
|
||||
Since `tidy3d`'s simulator registers itself as `jax`-differentiable (using the adjoint method), this "autodiff" support can extend all the way from parameters in the simulation definition, to gradients of the simulation output.
|
||||
When using these gradients for optimization, one achieves what is called "inverse design", where the desired properties of the output fields are used to automatically select simulation input parameters.
|
||||
|
||||
- **Performance**: XLA is a cross-industry project with the singular goal of providing a high-performance compilation target for data-driven programs.
|
||||
Published architects of OpenXLA include Alibaba, Amazon Web Services, AMD, Apple, Arm, Google, Intel, Meta, and NVIDIA.
|
||||
|
||||
- **Device Agnosticism**: XLA bytecode runs not just on CPUs, but on massively parallel devices like GPUs and TPUs as well.
|
||||
This enables massive speedups, and greatly expands the amount of data that is practical to work with at one time.
|
||||
|
||||
|
||||
Notes:
|
||||
The property `self.supports_jax` manually tracks whether these criteria are satisfied.
|
||||
|
||||
**As much as possible**, the _entirety of `blender_maxwell`_ is designed to maximize the ability to set `self.supports_jax = True` as often as possible.
|
||||
|
||||
**However**, there are many cases where a lazily-evaluated value is desirable, but `jax` isn't supported.
|
||||
These include design space exploration, where any particular parameter might vary for the purpose of producing batched simulations.
|
||||
In these cases, trying to compile a `self.func_jax` will raise a `ValueError`.
|
||||
|
||||
Returns:
|
||||
The `jit`-compiled function, ready to run on CPU, GPU, or XLA.
|
||||
|
||||
Raises:
|
||||
ValueError: If `self.supports_jax` is `False`.
|
||||
|
||||
References:
|
||||
JAX JIT: <https://jax.readthedocs.io/en/latest/jit-compilation.html>
|
||||
OpenXLA: <https://openxla.org/xla>
|
||||
"""
|
||||
if self.supports_jax:
|
||||
return jax.jit(self.func)
|
||||
|
||||
msg = 'Can\'t express LazyValueFuncFlow as JAX function (using jax.jit), since "self.supports_jax" is False'
|
||||
raise ValueError(msg)
|
||||
|
||||
####################
|
||||
# - Composition Operations
|
||||
####################
|
||||
def compose_within(
|
||||
self,
|
||||
enclosing_func: LazyFunction,
|
||||
enclosing_func_args: list[type] = (),
|
||||
enclosing_func_kwargs: dict[str, type] = MappingProxyType({}),
|
||||
supports_jax: bool = False,
|
||||
) -> typ.Self:
|
||||
"""Compose `self.func` within the given enclosing function, which itself takes arguments, and create a new `LazyValueFuncFlow` to contain it.
|
||||
|
||||
This is the fundamental operation used to "chain" functions together.
|
||||
|
||||
Examples:
|
||||
Consider a simple composition based on two expressions:
|
||||
```python
|
||||
R = spux.MathType.Real
|
||||
C = spux.MathType.Complex
|
||||
x, y = sp.symbols('x y', real=True)
|
||||
|
||||
# Prepare "Root" LazyValueFuncFlow w/x,y args
|
||||
expr_root = 3*x + y**2 - 100
|
||||
expr_root_func = sp.lambdify([x, y], expr, 'jax')
|
||||
|
||||
func_root = LazyValueFuncFlow(func=expr_root_func, func_args=[R,R], supports_jax=True)
|
||||
|
||||
# Compose "Enclosing" LazyValueFuncFlow w/z arg
|
||||
r = sp.Symbol('z', real=True)
|
||||
z = sp.Symbol('z', complex=True)
|
||||
expr = 10*sp.re(z) / (z + r)
|
||||
expr_func = sp.lambdify([r, z], expr, 'jax')
|
||||
|
||||
func = func_root.compose_within(enclosing_func=expr_func, enclosing_func_args=[C])
|
||||
|
||||
# Compute 'expr_func(expr_root_func(10.0, -500.0), 1+8j)'
|
||||
f.func_jax(10.0, -500.0, 1+8j)
|
||||
```
|
||||
|
||||
Using this function, it's easy to "keep adding" symbolic functions of any kind to the chain, without introducing extraneous complexity or compromising the ease of calling the final function.
|
||||
|
||||
Returns:
|
||||
A lazy function that takes both the enclosed and enclosing arguments, and returns the value of the enclosing function (whose first argument is the output value of the enclosed function).
|
||||
"""
|
||||
return LazyValueFuncFlow(
|
||||
func=lambda *args, **kwargs: enclosing_func(
|
||||
self.func(
|
||||
*list(args[: len(self.func_args)]),
|
||||
**{k: v for k, v in kwargs.items() if k in self.func_kwargs},
|
||||
),
|
||||
*args[len(self.func_args) :],
|
||||
**{k: v for k, v in kwargs.items() if k not in self.func_kwargs},
|
||||
),
|
||||
func_args=self.func_args + list(enclosing_func_args),
|
||||
func_kwargs=self.func_kwargs | dict(enclosing_func_kwargs),
|
||||
supports_jax=self.supports_jax and supports_jax,
|
||||
)
|
||||
|
||||
def __or__(
|
||||
self,
|
||||
other: typ.Self,
|
||||
):
|
||||
) -> typ.Self:
|
||||
"""Create a lazy function that takes all arguments of both lazy-function inputs, and itself promises to return a 2-tuple containing the outputs of both inputs.
|
||||
|
||||
Generally, `self.func` produces a single array as output (when doing math, at least).
|
||||
But sometimes (as in the `OperateMathNode`), we need to perform a binary operation between two arrays, like say, $+$.
|
||||
Without realizing both `LazyValueFuncFlow`s, it's not immediately obvious how one might accomplish this.
|
||||
|
||||
This overloaded function of the `|` operator (used as `left | right`) solves that problem.
|
||||
A new `LazyValueFuncFlow` is created, which takes the arguments of both inputs, and which produces a single output value: A 2-tuple, where each element if the output of each function.
|
||||
|
||||
Examples:
|
||||
Consider this illustrative (pseudocode) example:
|
||||
```python
|
||||
# Presume a,b are values, and that A,B are their identifiers.
|
||||
func_1 = LazyValueFuncFlow(func=compute_big_data_1, func_args=[A])
|
||||
func_2 = LazyValueFuncFlow(func=compute_big_data_2, func_args=[B])
|
||||
|
||||
f = (func_1 | func_2).compose_within(func=lambda D: D[0] + D[1])
|
||||
|
||||
f.func(a, b) ## Computes big_data_1 + big_data_2 @A=a, B=b
|
||||
```
|
||||
|
||||
Because of `__or__` (the operator `|`), the difficult and non-obvious task of adding the outputs of these unrealized functions because quite simple.
|
||||
|
||||
Notes:
|
||||
**Order matters**.
|
||||
`self` will be available in the new function's output as index `0`, while `other` will be available as index `1`.
|
||||
|
||||
As with anything lazy-composition-y, it can seem a bit strange at first.
|
||||
When reading the source code, pay special attention to the way that `args` is sliced to segment the positional arguments.
|
||||
|
||||
Returns:
|
||||
A lazy function that takes all arguments of both inputs, and returns a 2-tuple containing both output arguments.
|
||||
"""
|
||||
return LazyValueFuncFlow(
|
||||
func=lambda *args, **kwargs: (
|
||||
self.func(
|
||||
|
|
@ -178,33 +417,3 @@ class LazyValueFuncFlow:
|
|||
func_kwargs=self.func_kwargs | other.func_kwargs,
|
||||
supports_jax=self.supports_jax and other.supports_jax,
|
||||
)
|
||||
|
||||
# Composition
|
||||
def compose_within(
|
||||
self,
|
||||
enclosing_func: LazyFunction,
|
||||
enclosing_func_args: list[type] = (),
|
||||
enclosing_func_kwargs: dict[str, type] = MappingProxyType({}),
|
||||
supports_jax: bool = False,
|
||||
) -> typ.Self:
|
||||
return LazyValueFuncFlow(
|
||||
func=lambda *args, **kwargs: enclosing_func(
|
||||
self.func(
|
||||
*list(args[: len(self.func_args)]),
|
||||
**{k: v for k, v in kwargs.items() if k in self.func_kwargs},
|
||||
),
|
||||
*args[len(self.func_args) :],
|
||||
**{k: v for k, v in kwargs.items() if k not in self.func_kwargs},
|
||||
),
|
||||
func_args=self.func_args + list(enclosing_func_args),
|
||||
func_kwargs=self.func_kwargs | dict(enclosing_func_kwargs),
|
||||
supports_jax=self.supports_jax and supports_jax,
|
||||
)
|
||||
|
||||
@functools.cached_property
|
||||
def func_jax(self) -> LazyFunction:
|
||||
if self.supports_jax:
|
||||
return jax.jit(self.func)
|
||||
|
||||
msg = 'Can\'t express LazyValueFuncFlow as JAX function (using jax.jit), since "self.supports_jax" is False'
|
||||
raise ValueError(msg)
|
||||
|
|
|
|||
|
|
@ -24,9 +24,27 @@ import sympy as sp
|
|||
from blender_maxwell.utils import extra_sympy_units as spux
|
||||
from blender_maxwell.utils import logger
|
||||
|
||||
from .flow_kinds import FlowKind
|
||||
from .info import InfoFlow
|
||||
|
||||
log = logger.get(__name__)
|
||||
|
||||
|
||||
class ExprInfo(typ.TypedDict):
|
||||
active_kind: FlowKind
|
||||
size: spux.NumberSize1D
|
||||
mathtype: spux.MathType
|
||||
physical_type: spux.PhysicalType
|
||||
|
||||
# Value
|
||||
default_value: spux.SympyExpr
|
||||
|
||||
# LazyArrayRange
|
||||
default_min: spux.SympyExpr
|
||||
default_max: spux.SympyExpr
|
||||
default_steps: int
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True, kw_only=True)
|
||||
class ParamsFlow:
|
||||
func_args: list[spux.SympyExpr] = dataclasses.field(default_factory=list)
|
||||
|
|
@ -44,13 +62,24 @@ class ParamsFlow:
|
|||
return sorted(self.symbols, key=lambda sym: sym.name)
|
||||
|
||||
####################
|
||||
# - Scaled Func Args
|
||||
# - Realize Arguments
|
||||
####################
|
||||
def scaled_func_args(
|
||||
self,
|
||||
unit_system: spux.UnitSystem,
|
||||
symbol_values: dict[spux.Symbol, spux.SympyExpr] = MappingProxyType({}),
|
||||
):
|
||||
"""Realize the function arguments contained in this `ParamsFlow`, making it ready for insertion into `LazyValueFunc.func()`.
|
||||
|
||||
For all `arg`s in `self.func_args`, the following operations are performed:
|
||||
- **Unit System**: If `arg`
|
||||
|
||||
|
||||
Notes:
|
||||
This method is created for the purpose of being able to make this exact call in an `events.on_value_changed` method:
|
||||
|
||||
"""
|
||||
|
||||
"""Return the function arguments, scaled to the unit system, stripped of units, and cast to jax-compatible arguments."""
|
||||
if not all(sym in self.symbols for sym in symbol_values):
|
||||
msg = f"Symbols in {symbol_values} don't perfectly match the ParamsFlow symbols {self.symbols}"
|
||||
|
|
@ -112,3 +141,54 @@ class ParamsFlow:
|
|||
func_kwargs=self.func_kwargs | dict(enclosing_func_kwargs),
|
||||
symbols=self.symbols | enclosing_symbols,
|
||||
)
|
||||
|
||||
####################
|
||||
# - Generate ExprSocketDef
|
||||
####################
|
||||
def sym_expr_infos(
|
||||
self, info: InfoFlow, use_range: bool = False
|
||||
) -> dict[str, ExprInfo]:
|
||||
"""Generate all information needed to define expressions that realize all symbolic parameters in this `ParamsFlow`.
|
||||
|
||||
Many nodes need actual data, and as such, they require that the user select actual values for any symbols in the `ParamsFlow`.
|
||||
The best way to do this is to create one `ExprSocket` for each symbol that needs realizing.
|
||||
|
||||
Notes:
|
||||
This method is created for the purpose of being able to make this exact call in an `events.on_value_changed` method:
|
||||
```
|
||||
self.loose_input_sockets = {
|
||||
sym_name: sockets.ExprSocketDef(**expr_info)
|
||||
for sym_name, expr_info in params.sym_expr_infos(info).items()
|
||||
}
|
||||
```
|
||||
|
||||
Parameters:
|
||||
info: The InfoFlow associated with the `Expr` being realized.
|
||||
Each symbol in `self.symbols` **must** have an associated same-named dimension in `info`.
|
||||
use_range: Causes the
|
||||
|
||||
The `ExprInfo`s can be directly defererenced `**expr_info`)
|
||||
"""
|
||||
return {
|
||||
sym.name: {
|
||||
# Declare Kind/Size
|
||||
## -> Kind: Value prevents user-alteration of config.
|
||||
## -> Size: Always scalar, since symbols are scalar (for now).
|
||||
'active_kind': FlowKind.Value,
|
||||
'size': spux.NumberSize1D.Scalar,
|
||||
# Declare MathType/PhysicalType
|
||||
## -> MathType: Lookup symbol name in info dimensions.
|
||||
## -> PhysicalType: Same.
|
||||
'mathtype': info.dim_mathtypes[sym.name],
|
||||
'physical_type': info.dim_physical_types[sym.name],
|
||||
# TODO: Default Values
|
||||
# FlowKind.Value: Default Value
|
||||
#'default_value':
|
||||
# FlowKind.LazyArrayRange: Default Min/Max/Steps
|
||||
#'default_min':
|
||||
#'default_max':
|
||||
#'default_steps':
|
||||
}
|
||||
for sym in self.sorted_symbols
|
||||
if sym.name in info.dim_names
|
||||
}
|
||||
|
|
|
|||
|
|
@ -50,6 +50,7 @@ class NodeType(blender_type_enum.BlenderTypeEnum):
|
|||
# Outputs
|
||||
Viewer = enum.auto()
|
||||
## Outputs / File Exporters
|
||||
DataFileExporter = enum.auto()
|
||||
Tidy3DWebExporter = enum.auto()
|
||||
## Outputs / Web Exporters
|
||||
JSONFileExporter = enum.auto()
|
||||
|
|
|
|||
|
|
@ -19,13 +19,21 @@
|
|||
import dataclasses
|
||||
import enum
|
||||
import typing as typ
|
||||
from pathlib import Path
|
||||
|
||||
import jax.numpy as jnp
|
||||
import sympy as sp
|
||||
import jaxtyping as jtyp
|
||||
import numpy as np
|
||||
import polars as pl
|
||||
import tidy3d as td
|
||||
|
||||
from blender_maxwell.contracts import BLEnumElement
|
||||
from blender_maxwell.services import tdcloud
|
||||
from blender_maxwell.utils import extra_sympy_units as spux
|
||||
from blender_maxwell.utils import logger
|
||||
|
||||
from .flow_kinds.info import InfoFlow
|
||||
|
||||
log = logger.get(__name__)
|
||||
|
||||
|
||||
####################
|
||||
|
|
@ -293,3 +301,321 @@ class NewSimCloudTask:
|
|||
|
||||
task_name: tdcloud.CloudTaskName
|
||||
cloud_folder: tdcloud.CloudFolder
|
||||
|
||||
|
||||
####################
|
||||
# - Data File
|
||||
####################
|
||||
_DATA_FILE_EXTS = {
|
||||
'.txt',
|
||||
'.txt.gz',
|
||||
'.csv',
|
||||
'.npy',
|
||||
}
|
||||
|
||||
|
||||
class DataFileFormat(enum.StrEnum):
|
||||
"""Abstraction of a data file format, providing a regularized way of interacting with filesystem data.
|
||||
|
||||
Import/export interacts closely with the `Expr` socket's `FlowKind` semantics:
|
||||
- `FlowKind.LazyValueFunc`: Generally realized on-import/export.
|
||||
- **Import**: Loading data is generally eager, but memory-mapped file loading would be manageable using this interface.
|
||||
- **Export**: The function is realized and only the array is inserted into the file.
|
||||
- `FlowKind.Params`: Generally consumed.
|
||||
- **Import**: A new, empty `ParamsFlow` object is created.
|
||||
- **Export**: The `ParamsFlow` is consumed when realizing the `LazyValueFunc`.
|
||||
- `FlowKind.Info`: As the most important element, it is kept in an (optional) sidecar metadata file.
|
||||
- **Import**: The sidecar file is loaded, checked, and used, if it exists. A warning about further processing may show if it doesn't.
|
||||
- **Export**: The sidecar file is written next to the canonical data file, in such a manner that it can be both read and loaded.
|
||||
|
||||
Notes:
|
||||
This enum is UI Compatible, ex. for nodes/sockets desiring a dropdown menu of data file formats.
|
||||
|
||||
Attributes:
|
||||
Txt: Simple no-header text file.
|
||||
Only supports 1D/2D data.
|
||||
TxtGz: Identical to `Txt`, but compressed with `gzip`.
|
||||
Csv: Unspecific "Comma Separated Values".
|
||||
For loading, `pandas`-default semantics are used.
|
||||
For saving, very opinionated defaults are used.
|
||||
Customization is disabled on purpose.
|
||||
Npy: Generic numpy representation.
|
||||
Supports all kinds of numpy objects.
|
||||
Better laziness support via `jax`.
|
||||
"""
|
||||
|
||||
Csv = enum.auto()
|
||||
Npy = enum.auto()
|
||||
Txt = enum.auto()
|
||||
TxtGz = enum.auto()
|
||||
|
||||
####################
|
||||
# - UI
|
||||
####################
|
||||
@staticmethod
|
||||
def to_name(v: typ.Self) -> str:
|
||||
"""The extension name of the given `DataFileFormat`.
|
||||
|
||||
Notes:
|
||||
Called by the UI when creating an `EnumProperty` dropdown.
|
||||
"""
|
||||
return DataFileFormat(v).extension
|
||||
|
||||
@staticmethod
|
||||
def to_icon(v: typ.Self) -> str:
|
||||
"""No icon.
|
||||
|
||||
Notes:
|
||||
Called by the UI when creating an `EnumProperty` dropdown.
|
||||
"""
|
||||
return ''
|
||||
|
||||
def bl_enum_element(self, i: int) -> BLEnumElement:
|
||||
"""Produce a fully functional Blender enum element, given a particular integer index."""
|
||||
return (
|
||||
str(self),
|
||||
DataFileFormat.to_name(self),
|
||||
DataFileFormat.to_name(self),
|
||||
DataFileFormat.to_icon(self),
|
||||
i,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def bl_enum_elements() -> list[BLEnumElement]:
|
||||
"""Produce an immediately usable list of Blender enum elements, correctly indexed."""
|
||||
return [
|
||||
data_file_format.bl_enum_element(i)
|
||||
for i, data_file_format in enumerate(list(DataFileFormat))
|
||||
]
|
||||
|
||||
####################
|
||||
# - Properties
|
||||
####################
|
||||
@property
|
||||
def extension(self) -> str:
|
||||
"""Map to the actual string extension."""
|
||||
E = DataFileFormat
|
||||
return {
|
||||
E.Csv: '.csv',
|
||||
E.Npy: '.npy',
|
||||
E.Txt: '.txt',
|
||||
E.TxtGz: '.txt.gz',
|
||||
}[self]
|
||||
|
||||
####################
|
||||
# - Creation: Compatibility
|
||||
####################
|
||||
@staticmethod
|
||||
def valid_exts() -> list[str]:
|
||||
return _DATA_FILE_EXTS
|
||||
|
||||
@staticmethod
|
||||
def ext_has_valid_format(ext: str) -> bool:
|
||||
return ext in _DATA_FILE_EXTS
|
||||
|
||||
@staticmethod
|
||||
def path_has_valid_format(path: Path) -> bool:
|
||||
return path.is_file() and DataFileFormat.ext_has_valid_format(
|
||||
''.join(path.suffixes)
|
||||
)
|
||||
|
||||
def is_path_compatible(
|
||||
self, path: Path, must_exist: bool = False, can_exist: bool = True
|
||||
) -> bool:
|
||||
ext_matches = self.extension == ''.join(path.suffixes)
|
||||
match (must_exist, can_exist):
|
||||
case (False, False):
|
||||
return ext_matches and not path.is_file() and path.parent.is_dir()
|
||||
|
||||
case (True, False):
|
||||
msg = f'DataFileFormat: Path {path} cannot both be required to exist (must_exist=True), but also not be allowed to exist (can_exist=False)'
|
||||
raise ValueError(msg)
|
||||
|
||||
case (False, True):
|
||||
return ext_matches and path.parent.is_dir()
|
||||
|
||||
case (True, True):
|
||||
return ext_matches and path.is_file()
|
||||
|
||||
####################
|
||||
# - Creation
|
||||
####################
|
||||
@staticmethod
|
||||
def from_ext(ext: str) -> typ.Self | None:
|
||||
return {
|
||||
_ext: _data_file_ext
|
||||
for _data_file_ext, _ext in {
|
||||
k: k.extension for k in list(DataFileFormat)
|
||||
}.items()
|
||||
}.get(ext)
|
||||
|
||||
@staticmethod
|
||||
def from_path(path: Path) -> typ.Self | None:
|
||||
if DataFileFormat.path_has_valid_format(path):
|
||||
data_file_ext = DataFileFormat.from_ext(''.join(path.suffixes))
|
||||
if data_file_ext is not None:
|
||||
return data_file_ext
|
||||
|
||||
msg = f'DataFileFormat: Path "{path}" is compatible, but could not find valid extension'
|
||||
raise RuntimeError(msg)
|
||||
|
||||
return None
|
||||
|
||||
####################
|
||||
# - Functions: Metadata
|
||||
####################
|
||||
def supports_metadata(self) -> bool:
|
||||
E = DataFileFormat
|
||||
return {
|
||||
E.Csv: False, ## No RFC 4180 Support for Comments
|
||||
E.Npy: False, ## Quite simply no support
|
||||
E.Txt: True, ## Use # Comments
|
||||
E.TxtGz: True, ## Same as Txt
|
||||
}[self]
|
||||
|
||||
## TODO: Sidecar Metadata
|
||||
## - The vision is that 'saver' also writes metadata.
|
||||
## - This metadata is essentially a straight serialization of the InfoFlow.
|
||||
## - On-load, the metadata is used to re-generate the InfoFlow.
|
||||
## - This allows interpreting saved data without a ton of shenanigans.
|
||||
## - These sidecars could also be hand-writable for external data.
|
||||
## - When sidecars aren't found, the user would "fill in the blanks".
|
||||
## - ...Thus achieving the same result as if there were a sidecar.
|
||||
|
||||
####################
|
||||
# - Functions: Saver
|
||||
####################
|
||||
def is_info_compatible(self, info: InfoFlow) -> bool:
|
||||
E = DataFileFormat
|
||||
match self:
|
||||
case E.Csv:
|
||||
return len(info.dim_names) + info.output_shape_len <= 2
|
||||
case E.Npy:
|
||||
return True
|
||||
case E.Txt | E.TxtGz:
|
||||
return len(info.dim_names) + info.output_shape_len <= 2
|
||||
|
||||
@property
|
||||
def saver(
|
||||
self,
|
||||
) -> typ.Callable[[Path, jtyp.Shaped[jtyp.Array, '...'], InfoFlow], None]:
|
||||
def save_txt(path, data, info):
|
||||
np.savetxt(path, data)
|
||||
|
||||
def save_txt_gz(path, data, info):
|
||||
np.savetxt(path, data)
|
||||
|
||||
def save_csv(path, data, info):
|
||||
data_np = np.array(data)
|
||||
|
||||
# Extract Input Coordinates
|
||||
dim_columns = {
|
||||
dim_name: np.array(info.dim_idx_arrays[i])
|
||||
for i, dim_name in enumerate(info.dim_names)
|
||||
}
|
||||
|
||||
# Declare Function to Extract Output Values
|
||||
output_columns = {}
|
||||
|
||||
def declare_output_col(data_col, output_idx=0, use_output_idx=False):
|
||||
nonlocal output_columns
|
||||
|
||||
# Complex: Split to Two Columns
|
||||
output_idx_str = f'[{output_idx}]' if use_output_idx else ''
|
||||
if bool(np.any(np.iscomplex(data_col))):
|
||||
output_columns |= {
|
||||
f'{info.output_name}{output_idx_str}_re': np.real(data_col),
|
||||
f'{info.output_name}{output_idx_str}_im': np.imag(data_col),
|
||||
}
|
||||
|
||||
# Else: Use Array Directly
|
||||
else:
|
||||
output_columns |= {
|
||||
f'{info.output_name}{output_idx_str}': data_col,
|
||||
}
|
||||
|
||||
## TODO: Maybe a check to ensure dtype!=object?
|
||||
|
||||
# Extract Output Values
|
||||
## -> 2D: Iterate over columns by-index.
|
||||
## -> 1D: Declare the array as the only column.
|
||||
if len(data_np.shape) == 2:
|
||||
for output_idx in data_np.shape[1]:
|
||||
declare_output_col(data_np[:, output_idx], output_idx, True)
|
||||
else:
|
||||
declare_output_col(data_np)
|
||||
|
||||
# Compute DataFrame & Write CSV
|
||||
df = pl.DataFrame(dim_columns | output_columns)
|
||||
|
||||
log.debug('Writing Polars DataFrame to CSV:')
|
||||
log.debug(df)
|
||||
df.write_csv(path)
|
||||
|
||||
def save_npy(path, data, info):
|
||||
jnp.save(path, data)
|
||||
|
||||
E = DataFileFormat
|
||||
return {
|
||||
E.Csv: save_csv,
|
||||
E.Npy: save_npy,
|
||||
E.Txt: save_txt,
|
||||
E.TxtGz: save_txt_gz,
|
||||
}[self]
|
||||
|
||||
####################
|
||||
# - Functions: Loader
|
||||
####################
|
||||
@property
|
||||
def loader_is_jax_compatible(self) -> bool:
|
||||
E = DataFileFormat
|
||||
return {
|
||||
E.Csv: False,
|
||||
E.Npy: True,
|
||||
E.Txt: True,
|
||||
E.TxtGz: True,
|
||||
}[self]
|
||||
|
||||
@property
|
||||
def loader(
|
||||
self,
|
||||
) -> typ.Callable[[Path], tuple[jtyp.Shaped[jtyp.Array, '...'], InfoFlow]]:
|
||||
def load_txt(path: Path):
|
||||
return jnp.asarray(np.loadtxt(path))
|
||||
|
||||
def load_csv(path: Path):
|
||||
return jnp.asarray(pl.read_csv(path).to_numpy())
|
||||
## TODO: The very next Polars (0.20.27) has a '.to_jax' method!
|
||||
|
||||
def load_npy(path: Path):
|
||||
return jnp.load(path)
|
||||
|
||||
E = DataFileFormat
|
||||
return {
|
||||
E.Csv: load_csv,
|
||||
E.Npy: load_npy,
|
||||
E.Txt: load_txt,
|
||||
E.TxtGz: load_txt,
|
||||
}[self]
|
||||
|
||||
####################
|
||||
# - Metadata: Compatibility
|
||||
####################
|
||||
def is_info_compatible(self, info: InfoFlow) -> bool:
|
||||
E = DataFileFormat
|
||||
match self:
|
||||
case E.Csv:
|
||||
return len(info.dim_names) + (info.output_shape_len + 1) <= 2
|
||||
case E.Npy:
|
||||
return True
|
||||
case E.Txt | E.TxtGz:
|
||||
return len(info.dim_names) + (info.output_shape_len + 1) <= 2
|
||||
|
||||
def supports_metadata(self) -> bool:
|
||||
E = DataFileFormat
|
||||
return {
|
||||
E.Csv: False, ## No RFC 4180 Support for Comments
|
||||
E.Npy: False, ## Quite simply no support
|
||||
E.Txt: True, ## Use # Comments
|
||||
E.TxtGz: True, ## Same as Txt
|
||||
}[self]
|
||||
|
|
|
|||
|
|
@ -14,15 +14,10 @@
|
|||
# You should have received a copy of the GNU Affero General Public License
|
||||
# along with this program. If not, see <http://www.gnu.org/licenses/>.
|
||||
|
||||
import enum
|
||||
import typing as typ
|
||||
from pathlib import Path
|
||||
|
||||
import bpy
|
||||
import jax.numpy as jnp
|
||||
import jaxtyping as jtyp
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import sympy as sp
|
||||
import tidy3d as td
|
||||
|
||||
|
|
@ -35,112 +30,6 @@ from ... import base, events
|
|||
|
||||
log = logger.get(__name__)
|
||||
|
||||
####################
|
||||
# - Data File Extensions
|
||||
####################
|
||||
_DATA_FILE_EXTS = {
|
||||
'.txt',
|
||||
'.txt.gz',
|
||||
'.csv',
|
||||
'.npy',
|
||||
}
|
||||
|
||||
|
||||
class DataFileExt(enum.StrEnum):
|
||||
Txt = enum.auto()
|
||||
TxtGz = enum.auto()
|
||||
Csv = enum.auto()
|
||||
Npy = enum.auto()
|
||||
|
||||
####################
|
||||
# - Enum Elements
|
||||
####################
|
||||
@staticmethod
|
||||
def to_name(v: typ.Self) -> str:
|
||||
return DataFileExt(v).extension
|
||||
|
||||
@staticmethod
|
||||
def to_icon(v: typ.Self) -> str:
|
||||
return ''
|
||||
|
||||
####################
|
||||
# - Computed Properties
|
||||
####################
|
||||
@property
|
||||
def extension(self) -> str:
|
||||
"""Map to the actual string extension."""
|
||||
E = DataFileExt
|
||||
return {
|
||||
E.Txt: '.txt',
|
||||
E.TxtGz: '.txt.gz',
|
||||
E.Csv: '.csv',
|
||||
E.Npy: '.npy',
|
||||
}[self]
|
||||
|
||||
@property
|
||||
def loader(self) -> typ.Callable[[Path], jtyp.Shaped[jtyp.Array, '...']]:
|
||||
def load_txt(path: Path):
|
||||
return jnp.asarray(np.loadtxt(path))
|
||||
|
||||
def load_csv(path: Path):
|
||||
return jnp.asarray(pd.read_csv(path).values)
|
||||
|
||||
def load_npy(path: Path):
|
||||
return jnp.load(path)
|
||||
|
||||
E = DataFileExt
|
||||
return {
|
||||
E.Txt: load_txt,
|
||||
E.TxtGz: load_txt,
|
||||
E.Csv: load_csv,
|
||||
E.Npy: load_npy,
|
||||
}[self]
|
||||
|
||||
@property
|
||||
def loader_is_jax_compatible(self) -> bool:
|
||||
E = DataFileExt
|
||||
return {
|
||||
E.Txt: True,
|
||||
E.TxtGz: True,
|
||||
E.Csv: False,
|
||||
E.Npy: True,
|
||||
}[self]
|
||||
|
||||
####################
|
||||
# - Creation
|
||||
####################
|
||||
@staticmethod
|
||||
def from_ext(ext: str) -> typ.Self | None:
|
||||
return {
|
||||
_ext: _data_file_ext
|
||||
for _data_file_ext, _ext in {
|
||||
k: k.extension for k in list(DataFileExt)
|
||||
}.items()
|
||||
}.get(ext)
|
||||
|
||||
@staticmethod
|
||||
def from_path(path: Path) -> typ.Self | None:
|
||||
if DataFileExt.is_path_compatible(path):
|
||||
data_file_ext = DataFileExt.from_ext(''.join(path.suffixes))
|
||||
if data_file_ext is not None:
|
||||
return data_file_ext
|
||||
|
||||
msg = f'DataFileExt: Path "{path}" is compatible, but could not find valid extension'
|
||||
raise RuntimeError(msg)
|
||||
|
||||
return None
|
||||
|
||||
####################
|
||||
# - Compatibility
|
||||
####################
|
||||
@staticmethod
|
||||
def is_ext_compatible(ext: str):
|
||||
return ext in _DATA_FILE_EXTS
|
||||
|
||||
@staticmethod
|
||||
def is_path_compatible(path: Path):
|
||||
return path.is_file() and DataFileExt.is_ext_compatible(''.join(path.suffixes))
|
||||
|
||||
|
||||
####################
|
||||
# - Node
|
||||
|
|
@ -168,10 +57,6 @@ class DataFileImporterNode(base.MaxwellSimNode):
|
|||
def on_input_exprs_changed(self, input_sockets) -> None: # noqa: D102
|
||||
has_file_path = not ct.FlowSignal.check(input_sockets['File Path'])
|
||||
|
||||
has_file_path = ct.FlowSignal.check_single(
|
||||
input_sockets['File Path'], ct.FlowSignal.FlowPending
|
||||
)
|
||||
|
||||
if has_file_path:
|
||||
self.file_path = bl_cache.Signal.InvalidateCache
|
||||
|
||||
|
|
@ -188,10 +73,10 @@ class DataFileImporterNode(base.MaxwellSimNode):
|
|||
return None
|
||||
|
||||
@bl_cache.cached_bl_property(depends_on={'file_path'})
|
||||
def data_file_ext(self) -> DataFileExt | None:
|
||||
def data_file_format(self) -> ct.DataFileFormat | None:
|
||||
"""Retrieve the file extension by concatenating all suffixes."""
|
||||
if self.file_path is not None:
|
||||
return DataFileExt.from_path(self.file_path)
|
||||
return ct.DataFileFormat.from_path(self.file_path)
|
||||
return None
|
||||
|
||||
####################
|
||||
|
|
@ -201,7 +86,7 @@ class DataFileImporterNode(base.MaxwellSimNode):
|
|||
def expr_info(self) -> ct.InfoFlow | None:
|
||||
"""Retrieve the output expression's `InfoFlow`."""
|
||||
info = self.compute_output('Expr', kind=ct.FlowKind.Info)
|
||||
has_info = not ct.FlowKind.check(info)
|
||||
has_info = not ct.FlowSignal.check(info)
|
||||
if has_info:
|
||||
return info
|
||||
return None
|
||||
|
|
@ -216,13 +101,13 @@ class DataFileImporterNode(base.MaxwellSimNode):
|
|||
Called by Blender to determine the text to place in the node's header.
|
||||
"""
|
||||
if self.file_path is not None:
|
||||
return 'Load File: ' + self.file_path.name
|
||||
return 'Load: ' + self.file_path.name
|
||||
|
||||
return self.bl_label
|
||||
|
||||
def draw_info(self, _: bpy.types.Context, layout: bpy.types.UILayout) -> None:
|
||||
"""Show information about the loaded file."""
|
||||
if self.data_file_ext is not None:
|
||||
if self.data_file_format is not None:
|
||||
box = layout.box()
|
||||
row = box.row()
|
||||
row.alignment = 'CENTER'
|
||||
|
|
@ -235,16 +120,6 @@ class DataFileImporterNode(base.MaxwellSimNode):
|
|||
def draw_props(self, _: bpy.types.Context, layout: bpy.types.UILayout) -> None:
|
||||
pass
|
||||
|
||||
####################
|
||||
# - Events
|
||||
####################
|
||||
@events.on_value_changed(
|
||||
socket_name='File Path',
|
||||
input_sockets={'File Path'},
|
||||
)
|
||||
def on_file_changed(self, input_sockets) -> None:
|
||||
pass
|
||||
|
||||
####################
|
||||
# - FlowKind.Array|LazyValueFunc
|
||||
####################
|
||||
|
|
@ -264,19 +139,19 @@ class DataFileImporterNode(base.MaxwellSimNode):
|
|||
has_file_path = not ct.FlowSignal.check(input_sockets['File Path'])
|
||||
|
||||
if has_file_path:
|
||||
data_file_ext = DataFileExt.from_path(file_path)
|
||||
if data_file_ext is not None:
|
||||
data_file_format = ct.DataFileFormat.from_path(file_path)
|
||||
if data_file_format is not None:
|
||||
# Jax Compatibility: Lazy Data Loading
|
||||
## -> Delay loading of data from file as long as we can.
|
||||
if data_file_ext.loader_is_jax_compatible:
|
||||
if data_file_format.loader_is_jax_compatible:
|
||||
return ct.LazyValueFuncFlow(
|
||||
func=lambda: data_file_ext.loader(file_path),
|
||||
func=lambda: data_file_format.loader(file_path),
|
||||
supports_jax=True,
|
||||
)
|
||||
|
||||
# No Jax Compatibility: Eager Data Loading
|
||||
## -> Load the data now and bind it.
|
||||
data = data_file_ext.loader(file_path)
|
||||
data = data_file_format.loader(file_path)
|
||||
return ct.LazyValueFuncFlow(func=lambda: data, supports_jax=True)
|
||||
return ct.FlowSignal.FlowPending
|
||||
return ct.FlowSignal.FlowPending
|
||||
|
|
|
|||
|
|
@ -14,16 +14,15 @@
|
|||
# You should have received a copy of the GNU Affero General Public License
|
||||
# along with this program. If not, see <http://www.gnu.org/licenses/>.
|
||||
|
||||
# from . import file_exporters, viewer, web_exporters
|
||||
from . import viewer, web_exporters
|
||||
from . import file_exporters, viewer, web_exporters
|
||||
|
||||
BL_REGISTER = [
|
||||
*viewer.BL_REGISTER,
|
||||
# *file_exporters.BL_REGISTER,
|
||||
*file_exporters.BL_REGISTER,
|
||||
*web_exporters.BL_REGISTER,
|
||||
]
|
||||
BL_NODES = {
|
||||
**viewer.BL_NODES,
|
||||
# **file_exporters.BL_NODES,
|
||||
**file_exporters.BL_NODES,
|
||||
**web_exporters.BL_NODES,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -14,11 +14,15 @@
|
|||
# You should have received a copy of the GNU Affero General Public License
|
||||
# along with this program. If not, see <http://www.gnu.org/licenses/>.
|
||||
|
||||
from . import json_file_exporter
|
||||
from . import data_file_exporter
|
||||
|
||||
# from . import json_file_exporter
|
||||
|
||||
BL_REGISTER = [
|
||||
*json_file_exporter.BL_REGISTER,
|
||||
*data_file_exporter.BL_REGISTER,
|
||||
# *json_file_exporter.BL_REGISTER,
|
||||
]
|
||||
BL_NODES = {
|
||||
**json_file_exporter.BL_NODES,
|
||||
**data_file_exporter.BL_NODES,
|
||||
# **json_file_exporter.BL_NODES,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -0,0 +1,252 @@
|
|||
# blender_maxwell
|
||||
# Copyright (C) 2024 blender_maxwell Project Contributors
|
||||
#
|
||||
# This program is free software: you can redistribute it and/or modify
|
||||
# it under the terms of the GNU Affero General Public License as published by
|
||||
# the Free Software Foundation, either version 3 of the License, or
|
||||
# (at your option) any later version.
|
||||
#
|
||||
# This program is distributed in the hope that it will be useful,
|
||||
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
# GNU Affero General Public License for more details.
|
||||
#
|
||||
# You should have received a copy of the GNU Affero General Public License
|
||||
# along with this program. If not, see <http://www.gnu.org/licenses/>.
|
||||
|
||||
import typing as typ
|
||||
from pathlib import Path
|
||||
|
||||
import bpy
|
||||
|
||||
from blender_maxwell.utils import bl_cache, logger
|
||||
from blender_maxwell.utils import extra_sympy_units as spux
|
||||
|
||||
from .... import contracts as ct
|
||||
from .... import sockets
|
||||
from ... import base, events
|
||||
|
||||
log = logger.get(__name__)
|
||||
|
||||
|
||||
####################
|
||||
# - Operators
|
||||
####################
|
||||
class ExportDataFile(bpy.types.Operator):
|
||||
"""Exports data from the input to `DataFileExporterNode` to the file path given on the same node, if the path is compatible with the chosen export format (a property on the node)."""
|
||||
|
||||
bl_idname = ct.OperatorType.NodeExportDataFile
|
||||
bl_label = 'Save Data File'
|
||||
bl_description = 'Save a file with the contents, name, and format indicated by a NodeExportDataFile'
|
||||
|
||||
@classmethod
|
||||
def poll(cls, context):
|
||||
return (
|
||||
# Check Node
|
||||
hasattr(context, 'node')
|
||||
and hasattr(context.node, 'node_type')
|
||||
and (node := context.node).node_type == ct.NodeType.DataFileExporter
|
||||
# Check Expr
|
||||
and node.is_file_path_compatible_with_export_format
|
||||
)
|
||||
|
||||
def execute(self, context: bpy.types.Context):
|
||||
node = context.node
|
||||
|
||||
node.export_format.saver(node.file_path, node.expr_data, node.expr_info)
|
||||
return {'FINISHED'}
|
||||
|
||||
|
||||
####################
|
||||
# - Node
|
||||
####################
|
||||
class DataFileExporterNode(base.MaxwellSimNode):
|
||||
# """Export input data to a supported
|
||||
node_type = ct.NodeType.DataFileExporter
|
||||
bl_label = 'Data File Importer'
|
||||
|
||||
input_sockets: typ.ClassVar = {
|
||||
'Expr': sockets.ExprSocketDef(active_kind=ct.FlowKind.LazyValueFunc),
|
||||
'File Path': sockets.FilePathSocketDef(),
|
||||
}
|
||||
|
||||
####################
|
||||
# - Properties: Expr Info
|
||||
####################
|
||||
@events.on_value_changed(
|
||||
socket_name={'Expr'},
|
||||
input_sockets={'Expr'},
|
||||
input_socket_kinds={'Expr': ct.FlowKind.Info},
|
||||
)
|
||||
def on_input_exprs_changed(self, input_sockets) -> None: # noqa: D102
|
||||
has_expr = not ct.FlowSignal.check(input_sockets['Expr'])
|
||||
|
||||
if has_expr:
|
||||
self.expr_info = bl_cache.Signal.InvalidateCache
|
||||
|
||||
@bl_cache.cached_bl_property(depends_on={'file_path'})
|
||||
def expr_info(self) -> ct.InfoFlow | None:
|
||||
"""Retrieve the input expression's `InfoFlow`."""
|
||||
info = self._compute_input('Expr', kind=ct.FlowKind.Info)
|
||||
has_info = not ct.FlowSignal.check(info)
|
||||
if has_info:
|
||||
return info
|
||||
return None
|
||||
|
||||
@property
|
||||
def expr_data(self) -> typ.Any | None:
|
||||
"""Retrieve the input expression's data by evaluating its `LazyValueFunc`."""
|
||||
func = self._compute_input('Expr', kind=ct.FlowKind.LazyValueFunc)
|
||||
params = self._compute_input('Expr', kind=ct.FlowKind.Params)
|
||||
|
||||
has_func = not ct.FlowSignal.check(func)
|
||||
has_params = not ct.FlowSignal.check(params)
|
||||
if has_func and has_params:
|
||||
symbol_values = {
|
||||
sym.name: self._compute_input(sym.name, kind=ct.FlowKind.Value)
|
||||
for sym in params.sorted_symbols
|
||||
}
|
||||
return func.func_jax(
|
||||
*params.scaled_func_args(spux.UNITS_SI, symbol_values=symbol_values),
|
||||
**params.scaled_func_kwargs(spux.UNITS_SI, symbol_values=symbol_values),
|
||||
)
|
||||
|
||||
return None
|
||||
|
||||
####################
|
||||
# - Properties: File Path
|
||||
####################
|
||||
@events.on_value_changed(
|
||||
socket_name={'File Path'},
|
||||
input_sockets={'File Path'},
|
||||
input_socket_kinds={'File Path': ct.FlowKind.Value},
|
||||
input_sockets_optional={'File Path': True},
|
||||
)
|
||||
def on_file_path_changed(self, input_sockets) -> None: # noqa: D102
|
||||
has_file_path = not ct.FlowSignal.check(input_sockets['File Path'])
|
||||
if has_file_path:
|
||||
self.file_path = bl_cache.Signal.InvalidateCache
|
||||
|
||||
@bl_cache.cached_bl_property()
|
||||
def file_path(self) -> Path:
|
||||
"""Retrieve the input file path."""
|
||||
file_path = self._compute_input(
|
||||
'File Path', kind=ct.FlowKind.Value, optional=True
|
||||
)
|
||||
has_file_path = not ct.FlowSignal.check(file_path)
|
||||
if has_file_path:
|
||||
return file_path
|
||||
|
||||
return None
|
||||
|
||||
####################
|
||||
# - Properties: Export Format
|
||||
####################
|
||||
export_format: ct.DataFileFormat = bl_cache.BLField(
|
||||
enum_cb=lambda self, _: self.search_export_formats(),
|
||||
cb_depends_on={'expr_info'},
|
||||
)
|
||||
|
||||
def search_export_formats(self):
|
||||
if self.expr_info is not None:
|
||||
return [
|
||||
data_file_format.bl_enum_element(i)
|
||||
for i, data_file_format in enumerate(list(ct.DataFileFormat))
|
||||
if data_file_format.is_info_compatible(self.expr_info)
|
||||
]
|
||||
return ct.DataFileFormat.bl_enum_elements()
|
||||
|
||||
####################
|
||||
# - Properties: File Path Compatibility
|
||||
####################
|
||||
@bl_cache.cached_bl_property(depends_on={'file_path', 'export_format'})
|
||||
def is_file_path_compatible_with_export_format(self) -> bool | None:
|
||||
"""Determine whether the given file path is actually compatible with the desired export format."""
|
||||
if self.file_path is not None and self.export_format is not None:
|
||||
return self.export_format.is_path_compatible(self.file_path)
|
||||
return None
|
||||
|
||||
####################
|
||||
# - UI
|
||||
####################
|
||||
def draw_label(self):
|
||||
"""Show the extracted file name (w/extension) in the node's header label.
|
||||
|
||||
Notes:
|
||||
Called by Blender to determine the text to place in the node's header.
|
||||
"""
|
||||
if self.file_path is not None:
|
||||
return 'Save: ' + self.file_path.name
|
||||
|
||||
return self.bl_label
|
||||
|
||||
def draw_info(self, _: bpy.types.Context, layout: bpy.types.UILayout) -> None:
|
||||
"""Show information about the loaded file."""
|
||||
if self.export_format is not None:
|
||||
box = layout.box()
|
||||
row = box.row()
|
||||
row.alignment = 'CENTER'
|
||||
row.label(text='Data File')
|
||||
|
||||
row = box.row()
|
||||
row.alignment = 'CENTER'
|
||||
row.label(text=self.file_path.name)
|
||||
|
||||
compatibility = self.is_file_path_compatible_with_export_format
|
||||
if compatibility is not None:
|
||||
row = box.row()
|
||||
row.alignment = 'CENTER'
|
||||
if compatibility:
|
||||
row.label(text='Valid Path | Format', icon='CHECKMARK')
|
||||
else:
|
||||
row.label(text='Invalid Path | Format', icon='ERROR')
|
||||
|
||||
def draw_props(self, _: bpy.types.Context, layout: bpy.types.UILayout) -> None:
|
||||
layout.prop(self, self.blfields['export_format'], text='')
|
||||
layout.operator(ct.OperatorType.NodeExportDataFile, text='Save Data File')
|
||||
|
||||
####################
|
||||
# - Events
|
||||
####################
|
||||
@events.on_value_changed(
|
||||
# Trigger
|
||||
socket_name='Expr',
|
||||
run_on_init=True,
|
||||
# Loaded
|
||||
input_sockets={'Expr'},
|
||||
input_socket_kinds={'Expr': {ct.FlowKind.Info, ct.FlowKind.Params}},
|
||||
input_sockets_optional={'Expr': True},
|
||||
)
|
||||
def on_expr_changed(self, input_sockets: dict) -> None:
|
||||
"""Declare any loose input sockets needed to realize the input expr's symbols."""
|
||||
info = input_sockets['Expr'][ct.FlowKind.Info]
|
||||
params = input_sockets['Expr'][ct.FlowKind.Params]
|
||||
|
||||
has_info = not ct.FlowSignal.check(info)
|
||||
has_params = not ct.FlowSignal.check(params)
|
||||
|
||||
# Provide Sockets for Symbol Realization
|
||||
## -> Only happens if Params contains not-yet-realized symbols.
|
||||
if has_info and has_params and params.symbols:
|
||||
if set(self.loose_input_sockets) != {
|
||||
sym.name for sym in params.symbols if sym.name in info.dim_names
|
||||
}:
|
||||
self.loose_input_sockets = {
|
||||
sym_name: sockets.ExprSocketDef(**expr_info)
|
||||
for sym_name, expr_info in params.sym_expr_infos(info).items()
|
||||
}
|
||||
|
||||
elif self.loose_input_sockets:
|
||||
self.loose_input_sockets = {}
|
||||
|
||||
|
||||
####################
|
||||
# - Blender Registration
|
||||
####################
|
||||
BL_REGISTER = [
|
||||
ExportDataFile,
|
||||
DataFileExporterNode,
|
||||
]
|
||||
BL_NODES = {
|
||||
ct.NodeType.DataFileExporter: (ct.NodeCategory.MAXWELLSIM_OUTPUTS_FILEEXPORTERS)
|
||||
}
|
||||
Loading…
Reference in New Issue