RRAEsTorch 0.1.0__py3-none-any.whl

RRAEsTorch/wrappers/__init__.py

@@ -0,0 +1 @@
+from .wrappers import vmap_wrap, norm_wrap

RRAEsTorch/wrappers/wrappers.py

@@ -0,0 +1,237 @@
+import torch
+import dataclasses
+from dataclasses import dataclass
+from torch.func import vmap
+
+@dataclass(frozen=True)
+class NormParams:
+    pass
+
+@dataclass(frozen=True)
+class MeanStdParams(NormParams):
+    mean: float
+    std: float
+
+@dataclass(frozen=True)
+class MinMaxParams(NormParams):
+    min: float
+    max: float
+
+def find_norm_funcs(
+    array=None,
+    norm_typ="None",
+    params=None,
+    ):
+    """ Function that finds norm and inv_norm functions
+
+    The functions can be defined either based on an array (in_train),
+    In this case, the required parameters (e.g., mean, min, etc.) are
+    found using in_train. Otherwise, the parameters can be given
+    explicitly using params_in (check below for some examples).
+
+    Parameters
+    ----------
+    in_train : input array that is to be normalized
+    norm_in: Type of normalization, "minmax" and "meanstd" supported
+    params_in: Normalization parameters if they are to be given manually
+               instead of being computed from in_train.
+               These for example can be:
+               For mean/std normalization, params_out = {"mean": 0.12, "std": 0.23}
+               For min/max normalization, params_out = {"min": -1, "max": 3,}
+    Returns
+    -------
+        A new subclass of base_cls with methods that normalize the input when called.
+    """
+    if norm_typ != "None":
+        assert (params is not None) or (
+            array is not None
+        ), "Either params or in_train must be provided to set norm parameters"
+
+        assert not (
+            params is not None and array is not None
+        ), "Only one of params or in_train must be provided to set norm parameters"
+
+
+    match norm_typ:
+        case "minmax":
+            if params is None:
+                params = MinMaxParams(min=torch.min(array), max=torch.max(array))
+            else:
+                params = params
+            norm_fn = lambda self, x: (x - params.min) / (params.max - params.min)
+            inv_norm_fn = lambda self, x: x * (params.max - params.min) + params.min
+        case "meanstd":
+            if params is None:
+                params = MeanStdParams(mean=torch.mean(array), std=torch.std(array))
+            else:
+                params = params
+            norm_fn = lambda self, x: (x - params.mean) / params.std
+            inv_norm_fn = lambda self, x: x * params.std + params.mean
+        case "None":
+            if params is None:
+                params = NormParams()
+            else:
+                params = params
+            norm_fn = lambda self, x: x
+            inv_norm_fn = lambda self, x: x
+        case _:
+            raise NotImplementedError(f"norm_in specified {norm_typ} is not implemented.")
+
+    return norm_fn, inv_norm_fn, params
+
+def norm_in_wrap(base_cls, array=None, norm_typ="None", params=None, methods_to_wrap=["__call__"]):
+    """ Wrapper that normalizes the input of a function of a subclass of eqx.Module
+
+    The parameters of normalization can be either computed based on an array, or given
+    by the user.
+
+    Parameters
+    ----------
+    base_cls : Base class, subclass of eqx.Module of which functions will
+               be modified.
+    methods_to_wrap: Name of the methods in base_cls to be wrapped
+    array : The array from which norm parameters are to be found
+    norm_typ: Type of normalization, "meanstd" and "minmax" are supported
+    params: Parameters of normalization if these are to be given manually
+                These for example can be:
+                For mean/std normalization, params_out = {"mean": 0.12, "std": 0.23}
+                For min/max normalization, params_out = {"min": -1, "max": 3,}
+    Returns
+    -------
+        A new subclass of base_cls with methods that normalize the input when called.
+    """
+    norm_in, inv_norm_in, params_in = find_norm_funcs(array, norm_typ, params)
+    def norm_in_decorator(fn):
+        def wrapped(self, x, *args, **kwargs):
+            result = fn(self, norm_in(self, x), *args, **kwargs)
+            return result
+        return wrapped, {"norm_in": (callable, norm_in), "inv_norm_in": (callable, inv_norm_in), "params_in": (dict, params_in)}
+    return make_wrapped(base_cls, norm_in_decorator, methods_to_wrap)
+
+def inv_norm_out_wrap(base_cls, array=None, norm_typ="None", params=None, methods_to_wrap=["__call__"]):
+    """ Wrapper that de-normalizes the output of a function of a subclass of eqx.Module
+
+    The parameters of normalization can be either computed based on an array, or given
+    by the user.
+
+    Parameters
+    ----------
+    base_cls : Base class, subclass of eqx.Module of which functions will
+               be modified.
+    methods_to_wrap: Name of the methods in base_cls to be wrapped
+    array : The array from which norm parameters are to be found
+    norm_typ: Type of normalization, "meanstd" and "minmax" are supported
+    params: Parameters of normalization if these are to be given manually
+                These for example can be:
+                For mean/std normalization, params_out = {"mean": 0.12, "std": 0.23}
+                For min/max normalization, params_out = {"min": -1, "max": 3,}
+    Returns
+    -------
+        A new subclass of base_cls with de-normalized methods.
+        The given methods accept an additional keyword argument "keep_normalized"
+        If this argument is passed as True, functions retrieve original behavior
+    """
+    norm_out, inv_norm_out, params_out = find_norm_funcs(array, norm_typ, params)
+
+    def norm_out_decorator(fn):
+        def wrapped(self, x, *args, keep_normalized=False, **kwargs):
+            if keep_normalized:
+                result = fn(self, x, *args, **kwargs)
+            else:
+                result = inv_norm_out(self, fn(self, x, *args, **kwargs))
+            return result
+        return wrapped, {"norm_out": (callable, norm_out), "inv_norm_out": (callable, inv_norm_out), "params_out": (dict, params_out)}
+    return make_wrapped(base_cls, norm_out_decorator, methods_to_wrap)
+
+def norm_wrap(base_cls, array_in=None, norm_typ_in="None", params_in=None, array_out=None, norm_typ_out="None", params_out=None, methods_to_wrap_in=["__call__"], methods_to_wrap_out=["__call__"]):
+    """ Wrapper that normalizes functions of a subclass of eqx.Module
+
+    Parameters
+    ----------
+    base_cls : Base class, subclass of eqx.Module of which functions will
+               be modified.
+    methods_to_wrap: Name of the methods in base_cls to be wrapped
+    ... : Other parameters are explained in norm_in_wrap and norm_out_wrap
+
+    Returns
+    -------
+        A new subclass of base_cls with normalized methods.
+    """
+    after_in = norm_in_wrap(base_cls, array_in, norm_typ_in, params_in, methods_to_wrap_in)
+    after_out = inv_norm_out_wrap(after_in, array_out, norm_typ_out, params_out, methods_to_wrap_out)
+    return after_out
+
+def vmap_wrap(base_cls, map_axis, count=1, methods_to_wrap=["__call__"]):
+    """ Wrapper that vectorizes functions of a subclass of eqx.Module
+
+    Parameters
+    ----------
+    base_cls : Base class, subclass of eqx.Module of which functions will
+               be modified.
+    map_axis : Axis along which to vectorize the functions
+    methods_to_wrap: Name of the methods in base_cls to be wrapped
+    count: How many times to vectorize the functions
+
+    Returns
+    -------
+        A new subclass of base_cls with vectorized methods.
+        The given methods accept an additional keyword argument "no_map"
+        If this argument is passed as True, functions retrieve original behavior
+    """
+    def vmap_decorator(fn):
+        def wrapped(self, x, *args, no_map=False, **kwargs):
+            if (map_axis is None) or no_map:
+                return fn(self, x, *args, **kwargs)
+            f = lambda x: fn(self, x, *args, **kwargs)
+            for _ in range(count):
+                f = vmap(f, in_dims=(map_axis,), out_dims=map_axis)
+            out = f(x)
+            return out
+        return wrapped, {}
+    return make_wrapped(base_cls, vmap_decorator, methods_to_wrap)
+
+
+
+def make_wrapped(base_cls, decorator, methods_to_wrap=["__call__"]):
+    """
+    Create a subclass of base_cls with specified methods wrapped by decorator.
+
+    Parameters
+    ----------
+        base_cls: Original class to wrap.
+        methods_to_wrap: List of method names (strings) to decorate.
+        decorator: The wanted modification to the methods given above
+
+    Returns
+    -------
+        A new subclass of base_cls with decorated methods.
+    """
+    attrs = {}
+    annotations = {}
+    seen_fields = set()
+
+    for method_name in methods_to_wrap:
+        if not hasattr(base_cls, method_name):
+            raise AttributeError(f"Method {method_name} not found in {base_cls}")
+        
+        original = getattr(base_cls, method_name)
+        wrapped_method, extra_fields = decorator(original)
+        attrs[method_name] = wrapped_method
+
+        # Sort field names to avoid ordering issues
+        for field_name in extra_fields:  
+            if field_name in seen_fields:
+                continue
+
+            field_type, default_value = extra_fields[field_name]
+            annotations[field_name] = field_type
+
+            attrs[field_name] = dataclasses.field(default=default_value)
+
+            seen_fields.add(field_name)
+
+    if annotations:
+        attrs["__annotations__"] = annotations
+
+    Wrapped = type(base_cls.__name__, (base_cls,), attrs)
+    return Wrapped

rraestorch-0.1.0.dist-info/METADATA

@@ -0,0 +1,90 @@
+Metadata-Version: 2.4
+Name: RRAEsTorch
+Version: 0.1.0
+Summary:  A repo for RRAEs in PyTorch.
+Author-email: Jad Mounayer <jad.mounayer@outlook.com>
+License: MIT
+License-File: LICENSE
+License-File: LICENSE copy
+Requires-Python: >=3.10
+Requires-Dist: dill
+Requires-Dist: jaxtyping
+Requires-Dist: matplotlib
+Requires-Dist: numpy
+Requires-Dist: pandas
+Requires-Dist: prettytable
+Requires-Dist: pytest
+Requires-Dist: torch
+Requires-Dist: torchvision
+Requires-Dist: tqdm
+Description-Content-Type: text/markdown
+
+# RRAEsTorch
+
+A library that offers same functions as RRAEs (originally in JAX) but in PyTorch.
+
+The RAEDME is copied from RRAEs in the following:
+
+# Welcome
+
+This repository allows users to train and manipulate Equinox models easily, specifically, Autoencoders.
+
+The library provides trainor classes that allow to train Neural Networks in one line using JAX.
+
+It also provides easy ways to do normalization, and vectorization of matrices during training.
+
+There are also pre-built Autoencoder models, specifically Rank Reduction Autoencoders (RRAEs).
+
+# What are RRAEs?
+
+RRAEs or Rank reduction autoencoders are autoencoders include an SVD in the latent space to regularize the bottleneck.
+
+This library presents all the required classes for creating customized RRAEs and training them (other architectures such as Vanilla AEs, IRMAEs and LoRAEs are also available).
+
+Each script is an example of how to train a different model.
+
+To simply train an MLP (from equinox), try [this](general-MLP.py)
+
+To train an RRAE on curves (1D) using an MLP, refer to [this file](main-MLP.py)
+To train an RRAE on curves (1D) using an Convolutions, refer to [this file](main-CNN1D.py)
+To train an RRAE on images, refer to [this file](main-CNN.py)
+To train a VRRAE on images, refer to [this file](main-var-CNN.py)
+To train with an adaptive bottleneck size refer to [this](main-adap-MLP.py) and [this] file(main-adap-CNN.py)
+
+For examples of post-processing and what RRAE trainors can do, refer to [this file](post_process.py)
+
+# General instruction for preparing your own data
+
+In RRAEs.utilities, there's a function called get_data that can import many datasets to test.
+
+If you want to generate your own dataset, you will have to define the following:
+
+x_train: Train input (refer to each script to see the shape)
+
+x_test: Test input (refer to each script to see the shape)
+
+p_train: None (if you don't have any parameters, otherwise, 
+            these can be used for interpolation in the latent space)
+            
+p_test: Same as p_train
+
+y_train: = x_train for autoencoders
+
+y_test: = x_test for autoencoders
+
+pre_func_inp: lambda x:x (if not needed, this is a function to be applied
+                          on batches if memory is not enough to apply
+                          over whole dataset)
+                          
+pre_func_out: lambda x:x (same as above but for output)
+
+kwargs: {} (any other kwargs you might need)
+
+# Installation
+
+pip install RRAEsTorch
+
+Or to get the newest changes:
+
+pip install git+https://github.com/JadM133/RRAEsTorch.git
+

rraestorch-0.1.0.dist-info/RECORD

@@ -0,0 +1,27 @@
+RRAEsTorch/__init__.py,sha256=f234R6usRCqIgmBmiXyZNIHa7VrDe5E-KZO0Y6Ek5AQ,33
+RRAEsTorch/config.py,sha256=bQPwc_2KTvhglH_WIRSb5_6CpUQQj9AGpfqBp8_kuys,2931
+RRAEsTorch/AE_base/AE_base.py,sha256=4nnuv3VbrakkvpmcT9NqRBuFUZsC1BGMpQWdpxEcGyk,3398
+RRAEsTorch/AE_base/__init__.py,sha256=95YfMgEWzIFAkm--Ci-a9YPSGfCs2PDAK2sbfScT7oo,24
+RRAEsTorch/AE_classes/AE_classes.py,sha256=LTIKobJ5FXOPwGZLRhs82NKeuQzvXzj2YzdhuBLi2i0,18108
+RRAEsTorch/AE_classes/__init__.py,sha256=inM2_YPJG8T-lwx-CUg-zL2EMltmROQAlNZeZmnvVGA,27
+RRAEsTorch/tests/test_AE_classes_CNN.py,sha256=bEE9JnTo84t9w0a4kw1W74L51eLGjBB8trrlAG938RE,3182
+RRAEsTorch/tests/test_AE_classes_MLP.py,sha256=Cr1_uP7lag6RPQC1UhN2O7RFW4BEx1cd0Z-Y6VgrWRg,2718
+RRAEsTorch/tests/test_fitting_CNN.py,sha256=8i6oUZFS-DpaTh1VsRsd1rGG_Me7R3Kf1vO4lPSxYns,2797
+RRAEsTorch/tests/test_fitting_MLP.py,sha256=6Ggy5iAJmqMEJPDoon7fYkwFS8LWmosTOFBQuAQzszc,3353
+RRAEsTorch/tests/test_mains.py,sha256=ivTXP7NypSlgmB9VR5g0yq5VEuPZJGOibDqBMjOxHow,1021
+RRAEsTorch/tests/test_save.py,sha256=_zYcDZtz_HWzHHWT6uJVq7ynsFmlH5v6nDE-loZQ4zo,1997
+RRAEsTorch/tests/test_stable_SVD.py,sha256=OimHPqw4f22qndyRzwJfNvTzzjP2CM-yHtfXCqkMBuA,1230
+RRAEsTorch/tests/test_wrappers.py,sha256=Ike4IfMUx2Qic3f3_cBikgFPEU1WW5TuH1jT_r2NgvY,2215
+RRAEsTorch/trackers/__init__.py,sha256=3c9qcUMZiUfVr93rxFp6l11lIDthyK3PCY_-P-sNX3I,25
+RRAEsTorch/trackers/trackers.py,sha256=Pn1ejMxMjAtvgDazFFwa3qiZhogG5GtXj4UIIFiBpuY,9127
+RRAEsTorch/training_classes/__init__.py,sha256=K_Id4yhw640jp2JN15-0E4wJi4sPadi1fFRgovMV3kw,101
+RRAEsTorch/training_classes/training_classes.py,sha256=l1wPWZghYwzz_wxgrBi2fmOumgjbCz3qaaIcDks-GCg,36668
+RRAEsTorch/utilities/__init__.py,sha256=NtlizCcRW4qcsULXxWfjPk265rLJst0-GqWLRah2yDY,26
+RRAEsTorch/utilities/utilities.py,sha256=FzJWV9oFPF9sL9MC2m7euMqMKxCuLUEukzLfU0cF2to,53396
+RRAEsTorch/wrappers/__init__.py,sha256=txiLh4ylnuvPlapagz7DiAslmjllOzTqwCDL2dFr6dM,44
+RRAEsTorch/wrappers/wrappers.py,sha256=9Rmq2RS_EkZvsg96SKrt1HFIP35sF0xyPI0goV0ujOs,9659
+rraestorch-0.1.0.dist-info/METADATA,sha256=Wm_NiD_XGPLm8bHVI-U2Rt7aDnNAQ1tkoeeo3pzWarQ,3055
+rraestorch-0.1.0.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+rraestorch-0.1.0.dist-info/licenses/LICENSE,sha256=QQKXj7kEw2yUJxPDzHUnEPRO7YweIU6AAlwPLnhtinA,1090
+rraestorch-0.1.0.dist-info/licenses/LICENSE copy,sha256=QQKXj7kEw2yUJxPDzHUnEPRO7YweIU6AAlwPLnhtinA,1090
+rraestorch-0.1.0.dist-info/RECORD,,

rraestorch-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.28.0
+Root-Is-Purelib: true
+Tag: py3-none-any

rraestorch-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Jad Mounayer
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

rraestorch-0.1.0.dist-info/licenses/LICENSE copy

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Jad Mounayer
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.