[ ]:

#!pip install squarenet

[ ]:

from squarenet import SquareNet
from squarenet.sampler import samplepoints
from squarenet.backends import to_backend

Backend selection

[ ]:

backend = "jax"
#backend = "torch"

[8]:

if backend == "torch":
    import torch
if backend == "jax":
    import jax
    import jax.numpy as jnp

❒ SquareNet for PyTorch

I. sample a dataset

as a torch tensor, on the device of your choice

[9]:

I = 100
N = I*I*I
points = samplepoints("ball", size = (N, 3))
points = to_backend(points, backend = "torch")

type(points)

[9]:

torch.Tensor

Specify backend to the SquareNet

[10]:

sn = SquareNet(gridshape = (I, I, I), backend = "torch")
sn.fit(points)

Starting gridification... available method [fast, robust, ultimate]
selected fast

=== Carthesian sort ===
(max iter: 1000)
[█████████████████████████████] 1000/1000

=== Final Status ===
succesfully sorted at iteration 19

[11]:

grided_points = sn.map(points)
type(grided_points)

[11]:

torch.Tensor

[12]:

sn.plot(save = False)

../_images/examples_02_jax_and_pytorch_11_0.png

[12]:

(<Figure size 1200x400 with 3 Axes>, None)

❒ SquareNet for Jax

[14]:

I = 100
N = I*I
points = samplepoints("spiky", size = (N, 2))
points = to_backend(points, backend = "jax")

type(points)

[14]:

jaxlib._jax.ArrayImpl

[15]:

sn = SquareNet(gridshape = (I, I), backend = "jax")
sn.fit(points, method = "ultimate")

Starting gridification... available method [fast, robust, ultimate]
selected ultimate

=== Carthesian sort ===
(max iter: 2 x 1000 + 4 x 1000 + 1000)
[█████████████████████████████] 1000/1000
[█████████████████████████████] 1000/1000

[█████████████████████████████] 1000/1000
[█████████████████████████████] 1000/1000
[█████████████████████████████] 1000/1000
[█████████████████████████████] 1000/1000

[█████████████████████████████] 1000/1000

=== Final Status ===
succesfully sorted at iteration 95

[16]:

grided_points = sn.map(points)
type(grided_points)

[16]:

jaxlib._jax.ArrayImpl

[17]:

sn.plot(save = False)

../_images/examples_02_jax_and_pytorch_16_0.png

[17]:

(<Figure size 1200x400 with 3 Axes>, None)