python-RAT/ratapi/wrappers.py at df985a06306880ae3bcbb79fc4d7d87e8938fa5a · StephenNneji/python-RAT · GitHub

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
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
"""Wrappers for the interface between ratapi and MATLAB custom files."""

import os
import pathlib
from contextlib import suppress
from typing import Callable

import numpy as np
from numpy.typing import ArrayLike

import ratapi.rat_core


def start_matlab():
    """Start MATLAB asynchronously and returns a future to retrieve the engine later.

    Returns
    -------
    future : matlab.engine.futureresult.FutureResult
        A future used to get the actual matlab engine.

    """
    future = None
    if os.environ.get("DELAY_MATLAB_START", "0") == "0":
        with suppress(ImportError):
            import matlab.engine

            future = matlab.engine.start_matlab(background=True)

    return future


class MatlabWrapper:
    """Creates a python callback for a MATLAB function.

    Parameters
    ----------
    filename : string
        The path of the file containing MATLAB function

    """

    loader = start_matlab()
    loader_error_message = "matlabengine is required to use MatlabWrapper"

    def __init__(self, filename: str) -> None:
        if self.loader is None:
            raise ImportError(self.loader_error_message) from None

        self.engine = self.loader.result()
        path = pathlib.Path(filename)
        self.engine.cd(str(path.parent), nargout=0)
        self.function_name = path.stem

    def getHandle(self) -> Callable[[ArrayLike, ArrayLike, ArrayLike, int, int], tuple[ArrayLike, float]]:
        """Return a wrapper for the custom MATLAB function.

        Returns
        -------
        wrapper : Callable[[ArrayLike, ArrayLike, ArrayLike, int, int], tuple[ArrayLike, float]]
            The wrapper function for the MATLAB callback

        """

        def handle(*args):
            if len(args) == 2:
                output = getattr(self.engine, self.function_name)(
                    np.array(args[0], "float"),  # xdata
                    np.array(args[1], "float"),  # params
                    nargout=1,
                )
                return np.array(output, "float").tolist()
            else:
                matlab_args = [
                    np.array(args[0], "float"),  # params
                    np.array(args[1], "float"),  # bulk in
                    np.array(args[2], "float"),  # bulk out
                    float(args[3] + 1),  # contrast
                ]
                if len(args) > 4:
                    matlab_args.append(float(args[4] + 1))  # domain number

                output, sub_rough = getattr(self.engine, self.function_name)(
                    *matlab_args,
                    nargout=2,
                )
                return np.array(output, "float").tolist(), float(sub_rough)

        return handle


def use_shared_matlab(name, custom_error_message):
    """Connect asynchronously to shared MATLAB engine instance with the given name.

    Parameters
    ----------
    name : str
        The name of shared MATLAB engine instance
    custom_error_message : str
        The custom error message in case of failed connection

    Returns
    -------
    future : matlab.engine.futureresult.FutureResult
        A future used to get the actual matlab engine.

    """
    with suppress(ImportError):
        import matlab.engine

        MatlabWrapper.loader = matlab.engine.connect_matlab(name, background=True)
        MatlabWrapper.loader_error_message = custom_error_message
        return MatlabWrapper.loader


class DylibWrapper:
    """Creates a python callback for a function in dynamic library.

    Parameters
    ----------
    filename : str
        The path of the dynamic library
    function_name : str
        The name of the function to call

    """

    def __init__(self, filename, function_name) -> None:
        self.engine = ratapi.rat_core.DylibEngine(filename, function_name)

    def getHandle(self) -> Callable[[ArrayLike, ArrayLike, ArrayLike, int, int], tuple[ArrayLike, float]]:
        """Return a wrapper for the custom dynamic library function.

        Returns
        -------
        wrapper : Callable[[ArrayLike, ArrayLike, ArrayLike, int, int], tuple[ArrayLike, float]]
            The wrapper function for the dynamic library callback

        """

        def handle(*args):
            return self.engine.invoke(*args)

        return handle