Utilities¶
extract_paths(obj: PyTree, path='', op_type=None)
¶
Recursively extract paths to non-None leaves in a PyTree, including their operation type.
Source code in src/squint/utils/partition.py
partition_by_branches(pytree, branches_to_param)
¶
Partition a PyTree into parameters and static parts based on specified branch nodes.
All leaves that are descendants of any branch in branches_to_param
will be treated
as parameters; the rest will be considered static.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
pytree
|
PyTree
|
The input PyTree. |
required |
branches_to_param
|
list
|
A list of subtree objects whose leaves should be treated as parameters. |
required |
Returns:
Type | Description |
---|---|
(params_pytree, static_pytree) |
Source code in src/squint/utils/partition.py
partition_by_leaves(pytree, leaves_to_param)
¶
Partition a PyTree into parameters and static parts based on specified leaves. Args: pytree (PyTree): The input PyTree containing parameters and static parts. leaves_to_param (list): A list of leaves that should be treated as parameters. Returns: tuple: A tuple containing two PyTrees: the parameters and the static parts.
Example
import equinox as eqx leaves = [pytree.ops['phase'].phi, pytree.ops['phase'].epsilon] params, static = partition_by_leaves(pytree, leaves)
Source code in src/squint/utils/partition.py
partition_op(pytree: PyTree, name: Union[str, Sequence[str]])
¶
Partition a PyTree into parameters and static parts based on the operation name key. Args: pytree (PyTree): The input PyTree containing operations. name (str): The operation name key to filter by.
Source code in src/squint/utils/partition.py
print_nonzero_entries(arr)
¶
Print the indices and values of non-zero entries in a JAX array. Args: arr (jnp.ndarray): The JAX array to inspect.
Source code in src/squint/utils/__init__.py
hdfdict
¶
This code is adapted from, https://github.com/SiggiGue/hdfdict, which is licensed under MIT permissions.
LazyHdfDict
¶
Bases: UserDict
Helps loading data only if values from the dict are requested.
This is done by reimplementing the getitem method.
Source code in src/squint/utils/hdfdict.py
hdf_file(hdf, lazy=True, *args, **kwargs)
¶
Context manager that yields an h5 file if hdf
is a string,
otherwise it yields hdf as is.
Source code in src/squint/utils/hdfdict.py
unpack_dataset(item)
¶
Reconstruct a hdfdict dataset. Only some special unpacking for yaml and datetime types.
Parameters¶
item : h5py.Dataset
Returns¶
key: Unpacked key value : Unpacked Data
Source code in src/squint/utils/hdfdict.py
load(hdf, lazy=True, unpacker=unpack_dataset, mode='r', *args, **kwargs)
¶
Returns a dictionary containing the groups as keys and the datasets as values from given hdf file.
Parameters¶
hdf : string (path to file) or h5py.File()
or h5py.Group()
lazy : bool
If True, the datasets are lazy loaded at the moment an item is requested.
upacker : callable
Unpack function gets value
of type h5py.Dataset.
Must return the data you would like to have it in the returned dict.
mode : str
File read mode. Default: 'r'.
Returns¶
d : dict The dictionary containing all groupnames as keys and datasets as values, with group/file attributes under 'attrs'.
Source code in src/squint/utils/hdfdict.py
pack_dataset(hdfobject, key, value)
¶
Packs a given key value pair into a dataset in the given hdfobject.
Source code in src/squint/utils/hdfdict.py
dump(data, hdf, packer=pack_dataset, mode='w', *args, **kwargs)
¶
Adds keys of given dict as groups and values as datasets to the given hdf-file (by string or object) or group object.
Parameters¶
data : dict
The dictionary containing only string keys and
data values or dicts again.
hdf : string (path to file) or h5py.File()
or h5py.Group()
packer : callable
Callable gets hdfobject, key, value
as input.
hdfobject
is considered to be either a h5py.File or a h5py.Group.
key
is the name of the dataset.
value
is the dataset to be packed and accepted by h5py.
mode : str
File write mode. Default: 'w'
Returns¶
hdf : obj
h5py.Group()
or h5py.File()
instance
Source code in src/squint/utils/hdfdict.py
io
¶
IO
¶
The IO class encapsulates all saving/loading features of data, figures, etc. This provides consistent filetypes, naming conventions, etc.
Attributes:
Name | Type | Description |
---|---|---|
default_path |
Path
|
The default path where the data is stored. |
path |
Path
|
The path where the data is stored. |
verbose |
bool
|
A flag indicating whether to print out the path of each saved/loaded file. |
Typical usage
io = IO(path=r"\path\to\data") io.load_txt(filename="filename.txt")
or io = IO.create_new_save_folder(folder="subfolder", include_date=True, include_uuid=True) io.save_df(df, filename="dataframe.txt")
Source code in src/squint/utils/io.py
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 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 |
|
save_json(variable, filename)
¶
Save serialized python object into a json format, at filename
Parameters:
Name | Type | Description | Default |
---|---|---|---|
variable
|
The object to save. |
required | |
filename
|
str
|
Name of the file to which variable should be saved. |
required |
Returns:
Type | Description |
---|---|
None |
Source code in src/squint/utils/io.py
load_json(filename)
¶
Load serialized python object from json.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
filename
|
str
|
Name of the file from which we are loading the object. |
required |
Returns:
Type | Description |
---|---|
The loaded object data. |
Source code in src/squint/utils/io.py
save_txt(variable, filename)
¶
Save serialized python object into a text format, at filename.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
variable
|
The object to save. |
required | |
filename
|
str
|
Name of the file to which variable should be saved. |
required |
Returns:
Type | Description |
---|---|
None |
Source code in src/squint/utils/io.py
load_txt(filename)
¶
Load serialized python object from text file.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
filename
|
str
|
Name of the file from which we are loading the object. |
required |
Returns:
Type | Description |
---|---|
The loaded object data. |
Source code in src/squint/utils/io.py
save_dataframe(df, filename)
¶
Save a panda dataframe object to csv.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
df
|
DataFrame
|
Data contained in a dataframe. |
required |
filename
|
str
|
File to which data should be saved. |
required |
Returns:
Type | Description |
---|---|
None |
Source code in src/squint/utils/io.py
load_dataframe(filename)
¶
Load panda dataframe object from CSV.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
filename
|
str
|
Name of the file from which data should be loaded. |
required |
Returns:
Type | Description |
---|---|
pandas.DataFrame: Dataframe data. |
Source code in src/squint/utils/io.py
save_yaml(data, filename)
¶
Save dictionary to YAML file.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
filename
|
str
|
Name of the file from which data should be saved. |
required |
Source code in src/squint/utils/io.py
save_figure(fig, filename)
¶
Save a figure (image datatype can be specified as part of filename).
Parameters:
Name | Type | Description | Default |
---|---|---|---|
fig
|
Figure
|
The figure containing the figure to save. |
required |
filename
|
str
|
The filename to which we save a figure. |
required |
Returns:
Type | Description |
---|---|
None |
Source code in src/squint/utils/io.py
save_np_array(np_arr, filename)
¶
Save numpy array to a text document.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
np_arr
|
array
|
The array which we are saving. |
required |
filename
|
str
|
Name of the text file to which we want to save the numpy array. |
required |
Returns:
Type | Description |
---|---|
None |
Source code in src/squint/utils/io.py
load_np_array(filename, complex_vals=False)
¶
Loads numpy array from a text document.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
filename
|
str
|
Name of the text file from which we want to load the numpy array. |
required |
complex_vals
|
bool
|
True if we expect the numpy array to be complex, False otherwise. |
False
|
Returns:
Type | Description |
---|---|
numpy.array: The loaded numpy array. |
Source code in src/squint/utils/io.py
save_csv(df, filename)
¶
Save a panda dataframe object to csv.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
df
|
DataFrame
|
Data contained in a dataframe. |
required |
filename
|
str
|
File to which data should be saved. |
required |
Returns:
Type | Description |
---|---|
None |
Source code in src/squint/utils/io.py
load_csv(filename)
¶
Load panda dataframe object from CSV.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
filename
|
str
|
Name of the file from which data should be loaded. |
required |
Returns:
Type | Description |
---|---|
pandas.DataFrame: Dataframe data. |
Source code in src/squint/utils/io.py
save_h5(filename)
¶
Initialize an H5 file to save datasets into.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
filename
|
str
|
Name of the file from which data should be saved. |
required |
Returns:
Type | Description |
---|---|
h5py.File: H5 file. |
Source code in src/squint/utils/io.py
current_time()
¶
Returns the current date and time in a consistent format.
This function is used for monitoring long-running measurements by providing the current date and time in the "%d/%m/%Y, %H:%M:%S" format.
Returns:
Name | Type | Description |
---|---|---|
str |
The current date and time as a string in the "%d/%m/%Y, %H:%M:%S" format. |
Source code in src/squint/utils/io.py
partition
¶
partition_by_leaves(pytree, leaves_to_param)
¶
Partition a PyTree into parameters and static parts based on specified leaves. Args: pytree (PyTree): The input PyTree containing parameters and static parts. leaves_to_param (list): A list of leaves that should be treated as parameters. Returns: tuple: A tuple containing two PyTrees: the parameters and the static parts.
Example
import equinox as eqx leaves = [pytree.ops['phase'].phi, pytree.ops['phase'].epsilon] params, static = partition_by_leaves(pytree, leaves)
Source code in src/squint/utils/partition.py
partition_by_branches(pytree, branches_to_param)
¶
Partition a PyTree into parameters and static parts based on specified branch nodes.
All leaves that are descendants of any branch in branches_to_param
will be treated
as parameters; the rest will be considered static.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
pytree
|
PyTree
|
The input PyTree. |
required |
branches_to_param
|
list
|
A list of subtree objects whose leaves should be treated as parameters. |
required |
Returns:
Type | Description |
---|---|
(params_pytree, static_pytree) |
Source code in src/squint/utils/partition.py
partition_op(pytree: PyTree, name: Union[str, Sequence[str]])
¶
Partition a PyTree into parameters and static parts based on the operation name key. Args: pytree (PyTree): The input PyTree containing operations. name (str): The operation name key to filter by.
Source code in src/squint/utils/partition.py
extract_paths(obj: PyTree, path='', op_type=None)
¶
Recursively extract paths to non-None leaves in a PyTree, including their operation type.