1"""
  2# Description
  3
  4This module contains utility functions to handle multiple `qrotor.system` calculations.
  5These are commonly used as a list of `System` objects.
  6
  7
  8# Index
  9
 10| | |
 11| --- | --- |
 12| `as_list()`           | Ensures that a list only contains System objects |  
 13| `save_energies()`     | Save the energy eigenvalues for all systems to a CSV |  
 14| `save_splittings()`   | Save the tunnel splitting energies for all systems to a CSV |  
 15| `save_summary()`      | Save a summary of some relevant parameters for all systems to a CSV |  
 16| `get_energies()`      | Get the eigenvalues from all systems |  
 17| `get_gridsizes()`     | Get all gridsizes |  
 18| `get_runtimes()`      | Get all runtimes |  
 19| `get_tags()`          | Get a list with all system tags |  
 20| `filter_tags()`       | Filter the systems with or without specific tags |  
 21| `calculate_ideal_E()` | Calculate the ideal energy for a specified level |  
 22| `sort_by_gridsize()`  | Sort systems by gridsize |  
 23| `reduce_size()`       | Discard data that takes too much space |  
 24| `summary()`           | Print a summary of a System or list of Systems |  
 25
 26---
 27"""
 28
 29
 30from .system import System
 31from aton import txt
 32import pandas as pd
 33
 34
 35def as_list(systems) -> None:
 36    """Ensures that `systems` is a list of System objects.
 37
 38    If it is a System, returns a list with that System as the only element.
 39    If it is neither a list nor a System,
 40    or if the list does not contain only System objects,
 41    it raises an error.
 42    """
 43    if isinstance(systems, System):
 44        systems = [systems]
 45    if not isinstance(systems, list):
 46        raise TypeError(f"Must be a System object or a list of systems, found instead: {type(systems)}")
 47    for i in systems:
 48        if not isinstance(i, System):
 49            raise TypeError(f"All items in the list must be System objects, found instead: {type(i)}")
 50    return systems
 51
 52
 53def save_energies(
 54        systems:list,
 55        comment:str='',
 56        filepath:str='qrotor_eigenvalues.csv',
 57        ) -> pd.DataFrame:
 58    """Save the energy eigenvalues for all `systems` to a qrotor_eigenvalues.csv file.
 59
 60    Returns a Pandas Dataset with `System.comment` columns and `System.eigenvalues` values.
 61
 62    The output file can be changed with `filepath`,
 63    or set to null to avoid saving the dataset.
 64    A `comment` can be included at the top of the file.
 65    Note that `System.comment` must not include commas (`,`).
 66    """
 67    systems = as_list(systems)
 68    version = systems[0].version
 69    E = {}
 70    # Find max length of eigenvalues
 71    max_len = max((len(s.eigenvalues) if s.eigenvalues is not None else 0) for s in systems)
 72    for s in systems:
 73        if s.eigenvalues is not None:
 74            # Filter out None values and replace with NaN
 75            valid_eigenvalues = [float('nan') if e is None else e for e in s.eigenvalues]
 76            padded_eigenvalues = valid_eigenvalues + [float('nan')] * (max_len - len(s.eigenvalues))
 77        else:
 78            padded_eigenvalues = [float('nan')] * max_len
 79        E[s.comment] = padded_eigenvalues
 80    df = pd.DataFrame(E)
 81    if not filepath:
 82        return df
 83    # Else save to file
 84    df.to_csv(filepath, sep=',', index=False)
 85    # Include a comment at the top of the file
 86    file_comment = f'## {comment}\n' if comment else f''
 87    file_comment += f'# Energy eigenvalues\n'
 88    file_comment += f'# Calculated with QRotor {version}\n'
 89    file_comment += f'# https://pablogila.github.io/qrotor\n#'
 90    txt.edit.insert_at(filepath, file_comment, 0)
 91    print(f'Energy eigenvalues saved to {filepath}')
 92    return df
 93
 94
 95def save_splittings(
 96    systems:list,
 97    comment:str='',
 98    filepath:str='qrotor_splittings.csv',
 99    ) -> pd.DataFrame:
100    """Save the tunnel splitting energies for all `systems` to a qrotor_splittings.csv file.
101
102    Returns a Pandas Dataset with `System.comment` columns and `System.splittings` values.
103
104    The output file can be changed with `filepath`,
105    or set to null to avoid saving the dataset.
106    A `comment` can be included at the top of the file.
107    Note that `System.comment` must not include commas (`,`).
108    Different splitting lengths across systems are allowed - missing values will be NaN.
109    """
110    systems = as_list(systems)
111    version = systems[0].version
112    tunnelling_E = {}
113    # Find max length of splittings
114    max_len = max(len(s.splittings) for s in systems)
115    for s in systems:  # Pad shorter splittings with NaN
116        padded_splittings = s.splittings + [float('nan')] * (max_len - len(s.splittings))
117        tunnelling_E[s.comment] = padded_splittings
118    df = pd.DataFrame(tunnelling_E)
119    if not filepath:
120        return df
121    # Else save to file
122    df.to_csv(filepath, sep=',', index=False)
123    # Include a comment at the top of the file 
124    file_comment = f'## {comment}\n' if comment else f''
125    file_comment += f'# Tunnel splitting energies\n'
126    file_comment += f'# Calculated with QRotor {version}\n'
127    file_comment += f'# https://pablogila.github.io/qrotor\n#'
128    txt.edit.insert_at(filepath, file_comment, 0)
129    print(f'Tunnel splitting energies saved to {filepath}')
130    return df
131
132
133def save_summary(
134    systems:list,
135    comment:str='',
136    filepath:str='qrotor_summary.csv',
137    ) -> pd.DataFrame:
138    """Save a summary for all `systems` to a qrotor_summary.csv file.
139
140    Produces one row per System with the columns:
141    `comment`, `ZPE`, `E_activation`, `potential_max`, `1st_splitting`,
142    `1st_excitation`, `B`, `degeneracy`, `gridsize`.
143
144    Set `filepath` to null to just return the DataFrame.
145    """
146    systems = as_list(systems)
147    version = systems[0].version
148    rows = []
149    for s in systems:
150        eigenvalues = getattr(s, 'eigenvalues', None)
151        if eigenvalues is not None and len(eigenvalues) > 0:
152            first_val = eigenvalues[0]
153            zpe = float('nan') if first_val is None else first_val
154        else:
155            zpe = float('nan')
156        splittings = getattr(s, 'splittings', None)
157        if splittings is not None and len(splittings) > 0:
158            first_splitting = float('nan') if splittings[0] is None else splittings[0]
159        else:
160            first_splitting = float('nan')
161        excitations = getattr(s, 'excitations', None)
162        if excitations is not None and len(excitations) > 0:
163            first_excitation = float('nan') if excitations[0] is None else excitations[0]
164        else:
165            first_excitation = float('nan')
166        system_comment = getattr(s, 'comment', None)
167        E_activation = getattr(s, 'E_activation', None)
168        B = getattr(s, 'B', None)
169        tags = getattr(s, 'tags', None)
170        deg = getattr(s, 'deg', None)
171        gridsize = getattr(s, 'gridsize', None)
172        potential_max = getattr(s, 'potential_max', None)
173        # Each row contains the following:
174        rows.append({
175            'comment': system_comment,
176            'ZPE': zpe,
177            'potential_max': potential_max,
178            'E_activation': E_activation,
179            '1st_splitting': first_splitting,
180            '1st_excitation': first_excitation,
181            'B': B,
182            'degeneracy': deg,
183            'gridsize': gridsize,
184            'tags': tags,
185        })
186    # Save to file or just return df
187    df = pd.DataFrame(rows)
188    if not filepath:
189        return df
190    df.to_csv(filepath, sep=',', index=False)
191    # Include a comment at the top of the file
192    file_comment = f'## {comment}\n' if comment else ''
193    file_comment += '# Summary of systems\n'
194    file_comment += f'# Calculated with QRotor {version}\n'
195    file_comment += '# https://pablogila.github.io/qrotor\n#'
196    txt.edit.insert_at(filepath, file_comment, 0)
197    print(f'Summary saved to {filepath}')
198    return df
199
200
201def get_energies(systems:list) -> list:
202    """Get a list with all lists of eigenvalues from all systems.
203
204    If no eigenvalues are present for a particular system, appends None.
205    """
206    systems = as_list(systems)
207    energies = []
208    for i in systems:
209        if all(i.eigenvalues):
210            energies.append(i.eigenvalues)
211        else:
212            energies.append(None)
213    return energies
214
215
216def get_gridsizes(systems:list) -> list:
217    """Get a list with all gridsize values.
218
219    If no gridsize value is present for a particular system, appends None.
220    """
221    systems = as_list(systems)
222    gridsizes = []
223    for i in systems:
224        if i.gridsize:
225            gridsizes.append(i.gridsize)
226        elif any(i.potential_values):
227            gridsizes.append(len(i.potential_values))
228        else:
229            gridsizes.append(None)
230    return gridsizes
231
232
233def get_runtimes(systems:list) -> list:
234    """Returns a list with all runtime values.
235    
236    If no runtime value is present for a particular system, appends None.
237    """
238    systems = as_list(systems)
239    runtimes = []
240    for i in systems:
241        if i.runtime:
242            runtimes.append(i.runtime)
243        else:
244            runtimes.append(None)
245    return runtimes
246
247
248def get_tags(systems:list) -> list:
249    """Returns a list with all system tags."""
250    systems = as_list(systems)
251    tags = []
252    for i in systems:
253        # i.tags is guaranteed to exist and be a string (may be empty)
254        system_tags = i.tags.split()
255        for tag in system_tags:
256            if tag not in tags:
257                tags.append(tag)
258    return tags
259
260
261def filter_tags(
262        systems:list,
263        include:str='',
264        exclude:str='',
265        strict:bool=False,
266        ) -> list:
267    """Returns a filtered list of systems with or without specific tags.
268
269    You can `include` or `exclude` any number of tags, separated by blank spaces.
270    By default, the filters are triggered if any tag is found, i.e. *tag1 OR tag2*.
271    Set `strict=True` to require all tags to match, i.e. *tag1 AND tag2*.
272    """
273    systems = as_list(systems)
274    included_tags = include.split()
275    excluded_tags = exclude.split()
276    filtered_systems = []
277    for i in systems:
278        tags_found = get_tags(i)
279        if excluded_tags:
280            if strict and all(tag in tags_found for tag in excluded_tags):
281                continue
282            elif not strict and any(tag in tags_found for tag in excluded_tags):
283                continue
284        if included_tags:
285            if strict and not all(tag in tags_found for tag in included_tags):
286                continue
287            elif not strict and not any(tag in tags_found for tag in included_tags):
288                continue
289        filtered_systems.append(i)
290    return filtered_systems
291
292
293def calculate_ideal_E(E_level:int) -> int:
294    """Calculates the ideal energy for a specified `E_level`.
295
296    To be used in convergence tests with `potential_name = 'zero'`.
297    """
298    real_E_level = None
299    if E_level % 2 == 0:
300        real_E_level = E_level / 2
301    else:
302        real_E_level = (E_level + 1) / 2
303    ideal_E = int(real_E_level ** 2)
304    return ideal_E
305
306
307def sort_by_gridsize(systems:list) -> list:
308    """Sorts a list of System objects by `System.gridsize`."""
309    systems = as_list(systems)
310    systems = sorted(systems, key=lambda sys: sys.gridsize)
311    return systems
312
313
314def reduce_size(systems:list) -> list:
315    """Discard data that takes too much space.
316
317    Removes eigenvectors, potential values and grids,
318    for all System values inside the `systems` list.
319    """
320    systems = as_list(systems)
321    for dataset in systems:
322        dataset = dataset.reduce_size()
323    return systems
324
325
326def summary(
327        systems,
328        verbose:bool=False
329        ) -> None:
330    """Print a summary of a System or list of Systems.
331    
332    Print extra info with `verbose=True`
333    """
334    print('--------------------')
335    systems = as_list(systems)
336    
337    for system in systems:
338        dictionary = system.summary()
339        if verbose:
340            for key, value in dictionary.items():
341                print(f'{key:<24}', value)
342        else:
343            eigenvalues = system.eigenvalues if any(system.eigenvalues) else []
344            extra = ''
345            if len(system.eigenvalues) > 6:
346                eigenvalues = eigenvalues[:6]
347                extra = '...'
348            print('comment         ' + str(system.comment))
349            print('ZPE             ' + str(system.eigenvalues[0]))
350            print('E activation    ' + str(system.E_activation))
351            print('V max           ' + str(system.potential_max))
352            print('1st splitting   ' + str(system.splittings[0]))
353            print('1st excitation  ' + str(system.excitations[0]))
354            print('B               ' + str(system.B))
355            print('eigenvalues     ' + str([float(round(e, 4)) for e in eigenvalues]) + extra)
356            print('tags            ' + str(system.tags))
357            print('version         ' + str(system.version))
358        print('--------------------')
359    return None