diff --git a/tofu/imas2tofu/_core.py b/tofu/imas2tofu/_core.py index 39ee55ce1..4615d0180 100644 --- a/tofu/imas2tofu/_core.py +++ b/tofu/imas2tofu/_core.py @@ -614,7 +614,7 @@ def _rhosign(rho, theta): 'core_profiles':['t','Te','ne']} } - + _IDS_BASE = ['wall', 'pulse_schedule'] ################################### @@ -625,7 +625,8 @@ def _rhosign(rho, theta): def __init__(self, preset=None, dids=None, ids=None, occ=None, idd=None, shot=None, run=None, refshot=None, refrun=None, - user=None, tokamak=None, version=None, get=None, ref=True): + user=None, tokamak=None, version=None, + ids_base=None, synthdiag=None, get=None, ref=True): super(MultiIDSLoader, self).__init__() # Initialize dicts @@ -640,14 +641,22 @@ def __init__(self, preset=None, dids=None, ids=None, occ=None, idd=None, assert len(lidd) <= 1 idd = lidd[0] if len(lidd) > 0 else None self.add_ids(preset=preset, ids=ids, occ=occ, idd=idd, get=False) - if get is None and (ids is not None or preset is not None): + if ids_base is None: + ids_base = True + if ids_base is True: + self.add_ids_base(get=False) + if synthdiag is None: + synthdiag = False + if synthdiag is True: + self.add_ids_synthdiag(get=False) + if get is None and (len(self._dids) > 0 or preset is not None): get = True else: self.set_dids(dids) if get is None: get = True self._set_fsig() - if get: + if get is True: self.open_get_close() def _init_dict(self): @@ -1268,13 +1277,93 @@ def get_idd(self, idd=None): assert idd in self._didd.keys() return self._didd[idd]['idd'] + def _checkformat_ids_synthdiag(self, ids=None): + lc = [ids is None, isinstance(ids, str), isinstance(ids, list), + hasattr(ids, 'ids_properties')] + if not any(lc): + msg = ("Provided ids not understood!\n" + + "\t- provided: {}".format(str(ids))) + raise Exception(msg) + + lidssynth = [kk for kk, vv in self._didsdiag.items() + if 'synth' in vv.keys()] + if lc[0]: + ids = sorted(set(self._dids.keys()).intersection(lidssynth)) + elif lc[1]: + ids = [ids] + elif lc[3]: + ids = [ids.__class__.__name__] + + ids = sorted( + set(ids).intersection(lidssynth).intersection(self._dids.keys())) + if len(ids) == 0: + msg = ("The provided ids must be:\n" + + "\t- an is name (str)\n" + + "\t- a list of ids names\n" + + "\t- an ids instance\n" + + "\t- None\n" + + "And it must:\n" + + "\t- Already be added (cf. self.dids.keys())\n" + + "\t- Be a diagnostic ids with tabulated 'synth'") + # Turn to warning? => see user feedback + raise Exception(msg) + return ids + + def get_inputs_for_synthsignal(self, ids=None, verb=True, returnas=False): + """ Return and / or print a dict of the default inputs for desired ids + + Synthetic signal for a given diagnostic ids is computed from + signal that comes from other ids (e.g. core_profiles, equilibrium...) + For some diagnostics, the inputs required are already tabulated in + self._didsdiag[]['synth'] + + This method simply shows this already tabulated information + Advanced users may edit this hidden dictionnary to their needs + + """ + assert returnas in [False, True, dict, list] + ids = self._checkformat_ids_synthdiag(ids) + + # Deal with real case + if len(ids) == 1: + out = self._didsdiag[ids[0]]['synth'] + lids = sorted(out.get('dsig', {}).keys()) + if verb: + dmsg = ("\n\t-" + + "\n\t-".join([ + kk+':\n\t\t'+'\n\t\t'.join(vv) + for kk, vv in out.get('dsig', {}).items()])) + extra = {kk: vv for kk, vv in out.items() + if kk not in ['dsynth', 'dsig']} + msg = ("For computing synthetic signal for ids {}".format(ids) + + dmsg + '\n' + + "\t- Extra parameters (if any):\n" + + "\t\t{}\n".format(extra)) + print(msg) + if returnas is True: + returnas = dict + else: + out = None + lids = sorted(set(itt.chain.from_iterable([ + self._didsdiag[idsi]['synth'].get('dsig', {}).keys() + for idsi in ids]))) + if verb: + print(lids) + if returnas is True: + returnas = list + if returnas is dict: + return out + elif returnas is list: + return lids - def _checkformat_ids(self, ids, occ=None, idd=None, isget=None): + def _checkformat_ids(self, ids, occ=None, idd=None, isget=None, + synthdiag=False): # Check value and make dict if necessary lc = [type(ids) is str, type(ids) is list, - hasattr(ids, 'ids_properties')] + hasattr(ids, 'ids_properties'), + ids is None and synthdiag is True] if not any(lc): msg = "Arg ids must be either:\n" msg += " - str : valid ids name\n" @@ -1284,9 +1373,15 @@ def _checkformat_ids(self, ids, occ=None, idd=None, isget=None): msg += " Conditions: %s"%str(lc) raise Exception(msg) + # Synthdiag-specific + if synthdiag is True: + ids = self.get_inputs_for_synthsignal(ids=ids, verb=False, + returnas=list) + lc[1] = True + # Prepare dids[name] = {'ids':None/ids, 'needidd':bool} dids = {} - if lc[0]or lc[1]: + if lc[0] or lc[1]: if lc[0]: ids = [ids] for ids_ in ids: @@ -1323,7 +1418,6 @@ def _checkformat_ids(self, ids, occ=None, idd=None, isget=None): if dids[lids[ii]]['ids'] is not None: dids[lids[ii]]['ids'] = [dids[lids[ii]]['ids']]*nocc - # Format isget / get for ii in range(0,nids): nocc = dids[lids[ii]]['nocc'] @@ -1347,8 +1441,6 @@ def _checkformat_ids(self, ids, occ=None, idd=None, isget=None): return dids - - def add_ids(self, ids=None, occ=None, idd=None, preset=None, shot=None, run=None, refshot=None, refrun=None, user=None, tokamak=None, version=None, @@ -1395,14 +1487,44 @@ def add_ids(self, ids=None, occ=None, idd=None, preset=None, assert idd in self._didd.keys() # Add ids - if ids is not None: dids = self._checkformat_ids(ids, occ=occ, idd=idd, isget=isget) self._dids.update(dids) if get: - self.open_get_close(ids=ids) + self.open_get_close() + + def add_ids_base(self, occ=None, idd=None, + shot=None, run=None, refshot=None, refrun=None, + user=None, tokamak=None, version=None, + ref=None, isget=None, get=None): + """ Add th list of ids stored in self._IDS_BASE + + Typically used to add a list of common ids without having to re-type + them every time + """ + self.add_ids(ids=self._IDS_BASE, occ=occ, idd=idd, + shot=shot, run=run, refshot=refshot, refrun=refrun, + user=user, tokamak=tokamak, version=version, + ref=ref, isget=isget, get=get) + + def add_ids_for_synthdiag(self, ids=None, occ=None, idd=None, + shot=None, run=None, refshot=None, refrun=None, + user=None, tokamak=None, version=None, + ref=None, isget=None, get=None): + """ Add pre-tabulated input ids necessary for calculating synth. signal + The necessary input ids are given by self.get_inputs_for_synthsignal() + + """ + if get is None: + get = True + ids = self.get_inputs_for_synthsignal(ids=ids, verb=False, + returnas=list) + self.add_ids(ids=ids, occ=occ, idd=idd, preset=None, + shot=shot, run=run, refshot=refshot, refrun=refrun, + user=user, tokamak=tokamak, version=version, + ref=ref, isget=isget, get=get) def remove_ids(self, ids=None, occ=None): """ Remove an ids (optionally remove only an occurence) @@ -1529,11 +1651,12 @@ def __repr__(self): #--------------------- def _checkformat_getdata_ids(self, ids): - msg = "Arg ids must be either:\n" - msg += " - None: if self.dids only has one key\n" - msg += " - str: a valid key of self.dids\n\n" - msg += " Provided : %s\n"%ids - msg += " Available: %s"%str(list(self._dids.keys())) + msg = ("Arg ids must be either:\n" + + "\t- None: if self.dids only has one key\n" + + "\t- str: a valid key of self.dids\n\n" + + " Provided : {}\n".format(ids) + + " Available: {}\n".format(str(list(self._dids.keys()))) + + " => Consider using self.add_ids({})".format(str(ids))) lc = [ids is None, type(ids) is str] if not any(lc): @@ -1972,6 +2095,39 @@ def get_data_all(self, dsig=None, stack=True, warnings.warn(msg) return dout + def get_events(self, occ=None, verb=True, returnas=False): + """ Return chronoligical events stored in pulse_schedule + + If verb = True => print (default) + False => don't print + If returnas = list => return as list of tuples (name, time) + np.ndarray => return as np.ndarray + False => don't return (default) + """ + + # Check / format inputs + if verb is None: + verb = True + if returnas is None: + returnas = False + assert isinstance(verb, bool) + assert returnas in [False, list, tuple] + + events = self.get_data('pulse_schedule', + sig='events', occ=occ)['events'] + name, time = zip(*events) + ind = np.argsort(time) + if verb: + name, time = zip(*events[ind]) + msg = np.array([name, time], dtype='U').T + msg = np.char.ljust(msg, np.nanmax(np.char.str_len(msg))) + print(msg) + if returnas is list: + return events[ind].tolist() + elif returnas is tuple: + name, time = zip(*events[ind]) + return name, time + #--------------------- diff --git a/tofu/version.py b/tofu/version.py index ab0e46b95..52ca3f491 100644 --- a/tofu/version.py +++ b/tofu/version.py @@ -1,2 +1,2 @@ # Do not edit, pipeline versioning governed by git tags! -__version__ = '1.4.2-a5' +__version__ = '1.4.2-a5-15-g231dacd'