From 2502fa2a06ca45f5dbc817057c09948b82dd52c6 Mon Sep 17 00:00:00 2001 From: jgray-19 Date: Wed, 18 Oct 2023 01:32:39 +0200 Subject: [PATCH 01/11] Add first attempt at pandas dataframe --- src/pymadng/madp_classes.py | 44 +++++++++++++++++++++++++++++++++++++ 1 file changed, 44 insertions(+) diff --git a/src/pymadng/madp_classes.py b/src/pymadng/madp_classes.py index 29e2945..9ce9224 100644 --- a/src/pymadng/madp_classes.py +++ b/src/pymadng/madp_classes.py @@ -136,6 +136,50 @@ def __next__(self): except IndexError: raise StopIteration + def to_df(self): + """Converts the object to a pandas dataframe""" + if not self._mad.precv(f"MAD.typeid.is_mtable({self._name})"): + raise TypeError("Object is not a table, cannot convert to dataframe") + + import pandas as pd + py_name, name = self._mad.py_name, self._name + self._mad.send(f""" +-- Get the column names and column count +local colnames = {name}:colnames() +{py_name}:send({name}:ncol()) + +-- Loop through all the columns and send them +for i, colname in ipairs(colnames) do + local col = {name}:getcol(colname) + + -- If the column is a reference, convert it to a table + if not MAD.typeid.is_vector(col) or MAD.typeid.is_table(col) then + local tbl = table.new(#col, 0) + for i, val in ipairs(col) do tbl[i] = val end + col = tbl + end + + -- Send the column + {py_name}:send(colname):send(col) +end + +-- Get the header names and send the count +local header = {self._name}.header +{py_name}:send(#header) + +-- Loop through all the header names and send them +for i, attr in ipairs(header) do + {py_name}:send({self._name}[attr]):send(attr) +end +""") + ncol = self._mad.recv() + df = pd.DataFrame.from_dict( + {self._mad.recv(): np.array(self._mad.recv()).squeeze() for _ in range(ncol)} + ) + for _ in range(self._mad.recv()): + df.attrs[self._mad.recv()] = self._mad.recv() + return df + class madhl_fun(madhl_ref): # ----------------------------------Calling/Creating functions--------------------------------------# From 176cd7ceeccde59d7594f63a997c388cea245177 Mon Sep 17 00:00:00 2001 From: jgray-19 Date: Wed, 18 Oct 2023 01:33:37 +0200 Subject: [PATCH 02/11] Add contribution --- src/pymadng/madp_classes.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pymadng/madp_classes.py b/src/pymadng/madp_classes.py index 9ce9224..4925318 100644 --- a/src/pymadng/madp_classes.py +++ b/src/pymadng/madp_classes.py @@ -137,7 +137,7 @@ def __next__(self): raise StopIteration def to_df(self): - """Converts the object to a pandas dataframe""" + """Converts the object to a pandas dataframe, thanks to Pierre Schnizer for the idea and code.""" if not self._mad.precv(f"MAD.typeid.is_mtable({self._name})"): raise TypeError("Object is not a table, cannot convert to dataframe") From 9cf30bc80e7efe10435a829d1413fdaccac576f9 Mon Sep 17 00:00:00 2001 From: jgray-19 Date: Wed, 18 Oct 2023 01:34:10 +0200 Subject: [PATCH 03/11] Add testing --- examples/ex-ps-twiss/ps-twiss.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/examples/ex-ps-twiss/ps-twiss.py b/examples/ex-ps-twiss/ps-twiss.py index 832935f..b81a547 100644 --- a/examples/ex-ps-twiss/ps-twiss.py +++ b/examples/ex-ps-twiss/ps-twiss.py @@ -24,7 +24,10 @@ mad.py_strs_to_mad_strs(["name", "kind", "s", "l", "angle", "x", "y", "z", "theta"]), ) + import time + start = time.time() mad["mtbl", "mflw"] = mad.twiss(sequence=mad.ps, method=6, nslice=3, chrom=True) + print("twiss time:", time.time() - start) mad.load("MAD.gphys", "melmcol") #Add element properties as columns @@ -41,4 +44,10 @@ ).eval() #.eval() so tws:write() can be finished before MAD is shutdown + start = time.time() + df = mad.mtbl.to_df() + print("to_df time:", time.time() - start) + print(df) + print(df.attrs) + os.chdir(orginal_dir) \ No newline at end of file From 39e53001a735f1db34f66816e481a2d0fbc1fad6 Mon Sep 17 00:00:00 2001 From: jgray-19 Date: Wed, 18 Oct 2023 12:29:26 +0200 Subject: [PATCH 04/11] Add tfs-pandas, docstring and optional parameter columns --- examples/ex-ps-twiss/ps-twiss.py | 3 +- src/pymadng/madp_classes.py | 82 +++++++++++++++++++++----------- src/pymadng/madp_object.py | 4 +- 3 files changed, 57 insertions(+), 32 deletions(-) diff --git a/examples/ex-ps-twiss/ps-twiss.py b/examples/ex-ps-twiss/ps-twiss.py index b81a547..be9ae6c 100644 --- a/examples/ex-ps-twiss/ps-twiss.py +++ b/examples/ex-ps-twiss/ps-twiss.py @@ -1,4 +1,4 @@ -import time +import time, pandas from pymadng import MAD import numpy as np @@ -24,7 +24,6 @@ mad.py_strs_to_mad_strs(["name", "kind", "s", "l", "angle", "x", "y", "z", "theta"]), ) - import time start = time.time() mad["mtbl", "mflw"] = mad.twiss(sequence=mad.ps, method=6, nslice=3, chrom=True) print("twiss time:", time.time() - start) diff --git a/src/pymadng/madp_classes.py b/src/pymadng/madp_classes.py index 4925318..244fa1b 100644 --- a/src/pymadng/madp_classes.py +++ b/src/pymadng/madp_classes.py @@ -93,9 +93,7 @@ def __dir__(self) -> Iterable[str]: {self._mad.py_name}:send(modList) """ self._mad.psend(script) - varnames = [ - x for x in self._mad.recv() if isinstance(x, str) and x[0] != "_" - ] + varnames = [x for x in self._mad.recv() if isinstance(x, str) and x[0] != "_"] return varnames @@ -108,9 +106,7 @@ def __dir__(self) -> Iterable[str]: varnames = self._mad.precv(f"{self._name}:get_varkeys(MAD.object, false)") if not self._mad.ipython_use_jedi: - varnames.extend( - [x + "()" for x in self._mad.recv() if not x in varnames] - ) + varnames.extend([x + "()" for x in self._mad.recv() if not x in varnames]) return varnames def __call__(self, *args, **kwargs): @@ -136,20 +132,37 @@ def __next__(self): except IndexError: raise StopIteration - def to_df(self): - """Converts the object to a pandas dataframe, thanks to Pierre Schnizer for the idea and code.""" + def to_df(self, columns: list = None): + """Converts the object to a pandas dataframe, thanks to Pierre Schnizer for the idea and code. + + This function imports pandas and tfs-pandas, if tfs-pandas is not installed, it will only return a pandas dataframe. + + Args: + columns (list, optional): List of columns to include in the dataframe. Defaults to None. + + Returns: + pandas.DataFrame or tfs.TfsDataFrame: The dataframe containing the object's data. + """ if not self._mad.precv(f"MAD.typeid.is_mtable({self._name})"): raise TypeError("Object is not a table, cannot convert to dataframe") - + import pandas as pd + + try: + import tfs + + DataFrame, header = tfs.TfsDataFrame, "headers" + except ImportError: + DataFrame, header = pd.DataFrame, "attrs" + py_name, name = self._mad.py_name, self._name - self._mad.send(f""" + self._mad.send( + f""" -- Get the column names and column count -local colnames = {name}:colnames() {py_name}:send({name}:ncol()) --- Loop through all the columns and send them -for i, colname in ipairs(colnames) do +-- Loop through all the column names and send them with their data +for i, colname in ipairs({name}:colnames()) do local col = {name}:getcol(colname) -- If the column is a reference, convert it to a table @@ -159,7 +172,7 @@ def to_df(self): col = tbl end - -- Send the column + -- Send the column name and the column data {py_name}:send(colname):send(col) end @@ -169,17 +182,32 @@ def to_df(self): -- Loop through all the header names and send them for i, attr in ipairs(header) do - {py_name}:send({self._name}[attr]):send(attr) + {py_name}:send(attr):send({self._name}[attr]) end -""") - ncol = self._mad.recv() - df = pd.DataFrame.from_dict( - {self._mad.recv(): np.array(self._mad.recv()).squeeze() for _ in range(ncol)} +""" + ) + # Get the columns and header from the mad process + full_tbl = { # Not keen on the .squeeze() but it works (ng always sends 2D arrays, but I need the columns in 1D) + self._mad.recv(): np.array(self._mad.recv()).squeeze() + for _ in range(self._mad.recv()) # Evaluated first + } + + # Check if the user wants all the columns or just a few + if columns: + df = DataFrame.from_dict( + {key: val for key, val in full_tbl.items() if key in columns} ) - for _ in range(self._mad.recv()): - df.attrs[self._mad.recv()] = self._mad.recv() + else: + df = DataFrame.from_dict(full_tbl) + + # Get the header and add it to the dataframe + setattr( + df, + header, + {self._mad.recv(): self._mad.recv() for _ in range(self._mad.recv())}, + ) return df - + class madhl_fun(madhl_ref): # ----------------------------------Calling/Creating functions--------------------------------------# @@ -187,9 +215,7 @@ def __call_func(self, funcName: str, *args): """Call the function funcName and store the result in ``_last``.""" rtrn_ref = madhl_reflast(self._mad) args_string, vars_to_send = get_args_string(self._mad.py_name, *args) - self._mad.send( - f"{rtrn_ref._name} = __mklast__({funcName}({args_string}))\n" - ) + self._mad.send(f"{rtrn_ref._name} = __mklast__({funcName}({args_string}))\n") for var in vars_to_send: self._mad.send(var) return rtrn_ref @@ -225,9 +251,9 @@ class madhl_last: # The init and del for a _last object def __init__(self, mad_proc: mad_process): self._mad = mad_proc self._lst_cntr = mad_proc.lst_cntr - self._lastnum = mad_proc.lst_cntr.get() - self._name = f"_last[{self._lastnum}]" - self._parent = "_last" + self._lastnum = mad_proc.lst_cntr.get() + self._name = f"_last[{self._lastnum}]" + self._parent = "_last" def __del__(self): self._lst_cntr.set(self._lastnum) diff --git a/src/pymadng/madp_object.py b/src/pymadng/madp_object.py index 7f8c4d6..ebc0b7e 100644 --- a/src/pymadng/madp_object.py +++ b/src/pymadng/madp_object.py @@ -14,8 +14,8 @@ # TODO: Make it so that MAD does the loop for variables not python (speed) # TODO: Review recv_and exec: """ -Default arguments are evaluated once at module load time. -This may cause problems if the argument is a mutable object such as a list or a dictionary. +Default arguments are evaluated once at module load time. +This may cause problems if the argument is a mutable object such as a list or a dictionary. If the function modifies the object (e.g., by appending an item to a list), the default value is modified. Source: https://google.github.io/styleguide/pyguide.html """ From 66a76032649330e0bd9a175a64f35895805c5e6f Mon Sep 17 00:00:00 2001 From: jgray-19 Date: Wed, 18 Oct 2023 12:29:42 +0200 Subject: [PATCH 05/11] Add optional dependancies --- pyproject.toml | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 68fd3af..2949e89 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,3 +36,7 @@ where = ["src"] [tool.setuptools.dynamic] version = {attr = "pymadng.__version__"} + +[project.optional-dependencies] +pandas = ["pandas>=1.0,<2.1.0"] +tfs = ["tfs-pandas>3.0.0"] \ No newline at end of file From 462386ad1c4b65ff13d45c41b46539443de0e612 Mon Sep 17 00:00:00 2001 From: jgray-19 Date: Wed, 18 Oct 2023 17:42:38 +0200 Subject: [PATCH 06/11] Add a conversion to vector for speed. --- src/pymadng/madp_classes.py | 25 ++++++++++++++----------- 1 file changed, 14 insertions(+), 11 deletions(-) diff --git a/src/pymadng/madp_classes.py b/src/pymadng/madp_classes.py index 244fa1b..acae6da 100644 --- a/src/pymadng/madp_classes.py +++ b/src/pymadng/madp_classes.py @@ -87,14 +87,12 @@ def __dir__(self) -> Iterable[str]: name = self._name if name[:5] == "_last": name = name + ".__metatable or " + name - script = f""" + self._mad.psend(f""" local modList={{}}; local i = 1; for modname, mod in pairs({name}) do modList[i] = modname; i = i + 1; end {self._mad.py_name}:send(modList) - """ - self._mad.psend(script) - varnames = [x for x in self._mad.recv() if isinstance(x, str) and x[0] != "_"] - return varnames + """) + return [x for x in self._mad.recv() if isinstance(x, str) and x[0] != "_"] class madhl_obj(madhl_ref): @@ -133,7 +131,7 @@ def __next__(self): raise StopIteration def to_df(self, columns: list = None): - """Converts the object to a pandas dataframe, thanks to Pierre Schnizer for the idea and code. + """Converts the object to a pandas dataframe. This function imports pandas and tfs-pandas, if tfs-pandas is not installed, it will only return a pandas dataframe. @@ -156,7 +154,7 @@ def to_df(self, columns: list = None): DataFrame, header = pd.DataFrame, "attrs" py_name, name = self._mad.py_name, self._name - self._mad.send( + self._mad.send( # Sending every value individually is slow (sending vectors is fast) f""" -- Get the column names and column count {py_name}:send({name}:ncol()) @@ -165,11 +163,16 @@ def to_df(self, columns: list = None): for i, colname in ipairs({name}:colnames()) do local col = {name}:getcol(colname) - -- If the column is a reference, convert it to a table - if not MAD.typeid.is_vector(col) or MAD.typeid.is_table(col) then + -- If the column is not a vector and has a metatable, then convert it to a table (reference or generator columns) + if not MAD.typeid.is_vector(col) and getmetatable(col) then local tbl = table.new(#col, 0) - for i, val in ipairs(col) do tbl[i] = val end - col = tbl + conv_to_vec = true + for i, val in ipairs(col) do + tbl[i] = val + -- From testing, checking if I can convert to a vector is faster than sending the table + conv_to_vec = conv_to_vec and MAD.typeid.is_number(val) + end + col = conv_to_vec and MAD.vector(tbl) or tbl end -- Send the column name and the column data From c2bbe2754a75fb57b785b50ec95ea18c382d43d7 Mon Sep 17 00:00:00 2001 From: jgray-19 Date: Wed, 18 Oct 2023 17:44:46 +0200 Subject: [PATCH 07/11] Simplify creation and selection of dataframe --- src/pymadng/madp_classes.py | 18 ++++++------------ 1 file changed, 6 insertions(+), 12 deletions(-) diff --git a/src/pymadng/madp_classes.py b/src/pymadng/madp_classes.py index acae6da..a593102 100644 --- a/src/pymadng/madp_classes.py +++ b/src/pymadng/madp_classes.py @@ -189,24 +189,18 @@ def to_df(self, columns: list = None): end """ ) - # Get the columns and header from the mad process - full_tbl = { # Not keen on the .squeeze() but it works (ng always sends 2D arrays, but I need the columns in 1D) + # Create the dataframe from the data sent + # Not keen on the .squeeze() but it works (ng always sends 2D arrays, but I need the columns in 1D) + df = DataFrame.from_dict({ self._mad.recv(): np.array(self._mad.recv()).squeeze() for _ in range(self._mad.recv()) # Evaluated first - } + }) - # Check if the user wants all the columns or just a few if columns: - df = DataFrame.from_dict( - {key: val for key, val in full_tbl.items() if key in columns} - ) - else: - df = DataFrame.from_dict(full_tbl) + df = df[columns] # Only keep the columns specified # Get the header and add it to the dataframe - setattr( - df, - header, + setattr(df, header, {self._mad.recv(): self._mad.recv() for _ in range(self._mad.recv())}, ) return df From 4c878bd7c39b40e51a8bf83b849849249fc32bc2 Mon Sep 17 00:00:00 2001 From: jgray-19 Date: Wed, 18 Oct 2023 18:17:18 +0200 Subject: [PATCH 08/11] Simplify ps-twiss example --- examples/ex-ps-twiss/ps-twiss.py | 23 ++++++++++------------- 1 file changed, 10 insertions(+), 13 deletions(-) diff --git a/examples/ex-ps-twiss/ps-twiss.py b/examples/ex-ps-twiss/ps-twiss.py index be9ae6c..f542642 100644 --- a/examples/ex-ps-twiss/ps-twiss.py +++ b/examples/ex-ps-twiss/ps-twiss.py @@ -1,9 +1,6 @@ -import time, pandas +import os, time, pandas from pymadng import MAD -import numpy as np -import matplotlib.pyplot as plt -import os orginal_dir = os.getcwd() os.chdir(os.path.dirname(os.path.realpath(__file__))) @@ -24,9 +21,7 @@ mad.py_strs_to_mad_strs(["name", "kind", "s", "l", "angle", "x", "y", "z", "theta"]), ) - start = time.time() mad["mtbl", "mflw"] = mad.twiss(sequence=mad.ps, method=6, nslice=3, chrom=True) - print("twiss time:", time.time() - start) mad.load("MAD.gphys", "melmcol") #Add element properties as columns @@ -39,14 +34,16 @@ mad.mtbl.write("'PS_twiss_py.tfs'", mad.py_strs_to_mad_strs( ["name", "kind", "s", "x", "px", "beta11", "alfa11", "beta22", "alfa22","dx", - "dpx", "mu1", "mu2", "l", "angle", "k0l", "k1l", "k2l", "k3l", "hkick", "vkick"]), - ).eval() - #.eval() so tws:write() can be finished before MAD is shutdown - - start = time.time() + "dpx", "mu1", "mu2", "l", "angle", "k0l", "k1l", "k2l", "k3l", "hkick", "vkick"] + ) + ) + df = mad.mtbl.to_df() - print("to_df time:", time.time() - start) print(df) - print(df.attrs) + try: + import tfs + except ImportError: + print("tfs-pandas not installed, so the header is stored in attrs instead of headers") + print(df.attrs) os.chdir(orginal_dir) \ No newline at end of file From c3d42b8c860efb9df14f98f10b0aa40fc4ddbbf9 Mon Sep 17 00:00:00 2001 From: jgray-19 Date: Wed, 18 Oct 2023 18:59:39 +0200 Subject: [PATCH 09/11] Add documentation for dataframes --- docs/source/dataframes.rst | 12 ++++++++++++ docs/source/index.rst | 1 + docs/source/modules.rst | 14 +++++++++++--- docs/source/pymadng.rst | 10 ---------- examples/ex-ps-twiss/ps-twiss.py | 2 ++ 5 files changed, 26 insertions(+), 13 deletions(-) create mode 100644 docs/source/dataframes.rst delete mode 100644 docs/source/pymadng.rst diff --git a/docs/source/dataframes.rst b/docs/source/dataframes.rst new file mode 100644 index 0000000..a19cade --- /dev/null +++ b/docs/source/dataframes.rst @@ -0,0 +1,12 @@ +Converting TFS tables to Pandas DataFrames +------------------------------------------ + +The package `pandas` is an optional module, that has an inbuilt function to convert TFS tables (called ``mtable`` in MAD-NG) to a `pandas` ``DataFrame`` or a ``TfsDataFrame`` if you have `tfs-pandas` installed. In the example below, we generate an ``mtable`` by doing a survey and twiss on the Proton Synchrotron lattice, and then convert these to a ``DataFrame`` (or ``TfsDataFrame``). + +.. literalinclude:: ../../examples/ex-ps-twiss/ps-twiss.py + :lines: 18, 24, 41-49 + :linenos: + +In this script, we create the variables ``srv`` and ``mtbl`` which are ``mtable``s created by ``survey`` and ``twiss`` respectively. Then first, we convert the ``mtbl`` to a ``DataFrame`` and print it, before checking if you have `tfs-pandas` installed to check if we need to print out the header of the TFS table, which is stored in the attrs attribute of the ``DataFrame``, but is automatically printed when using `tfs-pandas`. Then we convert the ``srv`` to a ``DataFrame`` and print it. + +Note: If your object is not an ``mtable`` then this function will raise a ``TypeError``, but it is available to call on all ``object`` types in MAD-NG. \ No newline at end of file diff --git a/docs/source/index.rst b/docs/source/index.rst index 60a8b26..c26e13a 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -16,6 +16,7 @@ Welcome to the documentation for PyMAD-NG! ex-managing-refs ex-fodo ex-lhc-couplingLocal + dataframes ex-recv-lhc examples diff --git a/docs/source/modules.rst b/docs/source/modules.rst index 189f4d8..cb08903 100644 --- a/docs/source/modules.rst +++ b/docs/source/modules.rst @@ -1,7 +1,15 @@ API Reference ============= -.. toctree:: - :maxdepth: 4 +PyMAD-NG Module contents +------------------------ - pymadng +.. automodule:: pymadng + :members: + :undoc-members: + :show-inheritance: + +Useful functions for MAD References +----------------------------------- + +.. autofunction:: pymadng.madp_classes.madhl_obj.to_df \ No newline at end of file diff --git a/docs/source/pymadng.rst b/docs/source/pymadng.rst deleted file mode 100644 index 65abf9a..0000000 --- a/docs/source/pymadng.rst +++ /dev/null @@ -1,10 +0,0 @@ -PyMAD-NG -======== - -Module contents ---------------- - -.. automodule:: pymadng - :members: - :undoc-members: - :show-inheritance: diff --git a/examples/ex-ps-twiss/ps-twiss.py b/examples/ex-ps-twiss/ps-twiss.py index f542642..0875c91 100644 --- a/examples/ex-ps-twiss/ps-twiss.py +++ b/examples/ex-ps-twiss/ps-twiss.py @@ -46,4 +46,6 @@ print("tfs-pandas not installed, so the header is stored in attrs instead of headers") print(df.attrs) + print(mad.srv.to_df()) + os.chdir(orginal_dir) \ No newline at end of file From 80c1d9226e5eaca8852f915be1b74f09228fc831 Mon Sep 17 00:00:00 2001 From: jgray-19 Date: Wed, 18 Oct 2023 21:21:02 +0200 Subject: [PATCH 10/11] Add dataframe tests Make it stable for references --- src/pymadng/madp_classes.py | 37 ++++++++++------- tests/obj_tests.py | 80 ++++++++++++++++++++++++++++++++++++- 2 files changed, 101 insertions(+), 16 deletions(-) diff --git a/src/pymadng/madp_classes.py b/src/pymadng/madp_classes.py index a593102..1d6c4ad 100644 --- a/src/pymadng/madp_classes.py +++ b/src/pymadng/madp_classes.py @@ -153,15 +153,16 @@ def to_df(self, columns: list = None): except ImportError: DataFrame, header = pd.DataFrame, "attrs" - py_name, name = self._mad.py_name, self._name + py_name, obj_name = self._mad.py_name, self._name self._mad.send( # Sending every value individually is slow (sending vectors is fast) f""" --- Get the column names and column count -{py_name}:send({name}:ncol()) +-- Get the column names +colnames = {obj_name}:colnames() +{py_name}:send(colnames) -- Loop through all the column names and send them with their data -for i, colname in ipairs({name}:colnames()) do - local col = {name}:getcol(colname) +for i, colname in ipairs(colnames) do + local col = {obj_name}:getcol(colname) -- If the column is not a vector and has a metatable, then convert it to a table (reference or generator columns) if not MAD.typeid.is_vector(col) and getmetatable(col) then @@ -175,33 +176,39 @@ def to_df(self, columns: list = None): col = conv_to_vec and MAD.vector(tbl) or tbl end - -- Send the column name and the column data - {py_name}:send(colname):send(col) + -- Send the column data + {py_name}:send(col) end -- Get the header names and send the count -local header = {self._name}.header -{py_name}:send(#header) +local header = {obj_name}.header +{py_name}:send(header) -- Loop through all the header names and send them for i, attr in ipairs(header) do - {py_name}:send(attr):send({self._name}[attr]) + {py_name}:send({obj_name}[attr]) end """ ) # Create the dataframe from the data sent + colnames = self._mad.recv() + full_tbl = { # The string is in case references are within the table + col: self._mad.recv(f"{obj_name}:getcol('{col}')") for col in colnames + } + # Not keen on the .squeeze() but it works (ng always sends 2D arrays, but I need the columns in 1D) - df = DataFrame.from_dict({ - self._mad.recv(): np.array(self._mad.recv()).squeeze() - for _ in range(self._mad.recv()) # Evaluated first - }) + for key, val in full_tbl.items(): + if isinstance(val, np.ndarray): + full_tbl[key] = val.squeeze() + df = DataFrame(full_tbl) if columns: df = df[columns] # Only keep the columns specified # Get the header and add it to the dataframe + hnams = self._mad.recv() setattr(df, header, - {self._mad.recv(): self._mad.recv() for _ in range(self._mad.recv())}, + {hnam: self._mad.recv(f"{obj_name}['{hnam}']") for hnam in hnams} ) return df diff --git a/tests/obj_tests.py b/tests/obj_tests.py index 083cd24..c4b6789 100644 --- a/tests/obj_tests.py +++ b/tests/obj_tests.py @@ -1,4 +1,4 @@ -import unittest, os, time +import unittest, os, time, sys, tfs, pandas from pymadng import MAD from pymadng.madp_classes import madhl_ref, madhl_obj, madhl_fun @@ -299,6 +299,84 @@ def test_dir(self): self.assertEqual(dir(mad.quadrupole(knl=[0, 0.3], l = 1)), quad_exp) #Dir of instance of class should be the same as the class self.assertEqual(dir(mad.quadrupole(asd = 10, qwe = 20)), sorted(quad_exp + ["asd", "qwe"])) #Adding to the instance should change the dir +class TestDataFrame(unittest.TestCase): + + def generalDataFrame(self, headers, DataFrame): + mad = MAD() + mad.send(""" +test = mtable{ + {"string"}, "number", "integer", "complex", "boolean", "list", "table", "range",! "generator", + name = "test", + header = {"string", "number", "integer", "complex", "boolean", "list", "table", "range"}, + string = "string", + number = 1.234567890, + integer = 12345670, + complex = 1.3 + 1.2i, + boolean = true, + list = {1, 2, 3, 4, 5}, + table = {1, 2, ["key"] = "value"}, + range = 1..11, +} + + {"a", 1.1, 1, 1 + 2i, true , {1, 2 }, {1 , 2 , ["3" ] = 3 }, 1..11,} + + {"b", 2.2, 2, 2 + 3i, false, {3, 4 }, {4 , 5 , ["6" ] = 6 }, 2..12,} + + {"c", 3.3, 3, 3 + 4i, true , {5, 6 }, {7 , 8 , ["9" ] = 9 }, 3..13,} + + {"d", 4.4, 4, 4 + 5i, false, {7, 8 }, {10, 11, ["12"] = 12}, 4..14,} + + {"e", 5.5, 5, 5 + 6i, true , {9, 10}, {13, 14, ["15"] = 15}, 5..15,} + +test:addcol("generator", \\ri, m -> m:getcol("number")[ri] + 1i * m:getcol("number")[ri]) +test:write("test") + """ + ) + df = mad.test.to_df() + self.assertTrue(isinstance(df, DataFrame)) + self.assertEqual(getattr(df, headers)["name"], "test") + self.assertEqual(getattr(df, headers)["string"], "string") + self.assertEqual(getattr(df, headers)["number"], 1.234567890) + self.assertEqual(getattr(df, headers)["integer"], 12345670) + self.assertEqual(getattr(df, headers)["complex"], 1.3 + 1.2j) + self.assertEqual(getattr(df, headers)["boolean"], True) + self.assertEqual(getattr(df, headers)["list"], [1, 2, 3, 4, 5]) + lst, hsh = getattr(df, headers)["table"] + self.assertEqual(lst, [1, 2]) + self.assertEqual(hsh["key"], "value") + + self.assertEqual(df["string"].tolist(), ["a", "b", "c", "d", "e"]) + self.assertEqual(df["number"].tolist(), [1.1, 2.2, 3.3, 4.4, 5.5]) + self.assertEqual(df["integer"].tolist(), [1, 2, 3, 4, 5]) + self.assertEqual(df["complex"].tolist(), [1 + 2j, 2 + 3j, 3 + 4j, 4 + 5j, 5 + 6j]) + self.assertEqual(df["boolean"].tolist(), [True, False, True, False, True]) + self.assertEqual(df["list"].tolist(), [[1, 2], [3, 4], [5, 6], [7, 8], [9, 10]]) + tbl = df["table"].tolist() + for i in range(len(tbl)): + lst, hsh = tbl[i] + self.assertEqual(lst, [i*3 + 1, i*3 + 2]) + self.assertEqual(hsh[str((i+1) * 3)], (i+1) * 3) + self.assertEqual( + df["range"].tolist(), + [range(1, 12), range(2, 13), range(3, 14), range(4, 15), range(5, 16)] + ) + + def testTfsDataFrame(self): + self.generalDataFrame("headers", tfs.TfsDataFrame) + + def testPandasDataFrame(self): + sys.modules["tfs"] = None #Remove tfs-pandas + self.generalDataFrame("attrs", pandas.DataFrame) + del sys.modules["tfs"] + + def testFailure(self): + with MAD() as mad: + mad.send(""" +test = mtable{"string", "number"} + {"a", 1.1} + {"b", 2.2} + """) + pandas = sys.modules["pandas"] + sys.modules["pandas"] = None + self.assertRaises(ImportError, lambda: mad.test.to_df()) + sys.modules["pandas"] = pandas + df = mad.test.to_df() + self.assertTrue(isinstance(df, tfs.TfsDataFrame)) + self.assertEqual(df["string"].tolist(), ["a", "b"]) + self.assertEqual(df["number"].tolist(), [1.1, 2.2]) class TestSpeed(unittest.TestCase): From 390fc13a01088761f3b83da5eb258bab2c5b9c68 Mon Sep 17 00:00:00 2001 From: jgray-19 Date: Wed, 18 Oct 2023 21:31:24 +0200 Subject: [PATCH 11/11] Update Version --- CHANGELOG.md | 4 ++++ src/pymadng/__init__.py | 2 +- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index d88889e..214f1b0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,7 @@ +0.4.2 (2023/10/18) + +Add `to_df` method to objects, allowing for easy conversion to pandas dataframes. \ + 0.4.1 (2023/08/19) Change the way `send_vars` and `recv_vars` work, they now use kwargs and args respectively. \ diff --git a/src/pymadng/__init__.py b/src/pymadng/__init__.py index 4ee3414..101d324 100644 --- a/src/pymadng/__init__.py +++ b/src/pymadng/__init__.py @@ -1,7 +1,7 @@ from .madp_object import MAD __title__ = "pymadng" -__version__ = "0.4.1" +__version__ = "0.4.2" __summary__ = "Python interface to MAD-NG running as subprocess" __uri__ = "https://github.com/MethodicalAcceleratorDesign/MADpy"