diff --git a/tofu/data/_comp.py b/tofu/data/_comp.py index a18728a1b..77abb52cb 100644 --- a/tofu/data/_comp.py +++ b/tofu/data/_comp.py @@ -450,8 +450,8 @@ def func(pts, vect=None, t=None, ntall=ntall, indpts = trifind(r,z) indok = indpts > -1 if t is None: - for ii in range(0,ntall): - val[ii,indok] = vquant[indtq[ii],indpts[indok]] + for ii in range(0, ntall): + val[ii, indok] = vquant[indtq[ii], indpts[indok]] t = tall else: ntall, indt, indtu = plasma._get_indtu(t=t, tall=tall, diff --git a/tofu/data/_core.py b/tofu/data/_core.py index adbfd68e7..8913c41cd 100644 --- a/tofu/data/_core.py +++ b/tofu/data/_core.py @@ -2371,9 +2371,9 @@ def _checkformat_dtrm(dtime=None, dradius=None, dmesh=None, # Define allowed keys for each dict lkok = ['data', 'dim', 'quant', 'name', 'origin', 'units', 'depend'] - lkmeshmax = ['type','ftype','nodes','faces', - 'nfaces','nnodes','mpltri','size','ntri'] - lkmeshmin = ['type','ftype','nodes','faces'] + lkmeshmax = ['type', 'ftype', 'nodes', 'faces', 'R', 'Z', 'shapeRZ', + 'nfaces', 'nnodes', 'mpltri', 'size', 'ntri'] + lkmeshmin = ['type', 'ftype'] dkok = {'dtime': {'max':lkok, 'min':['data'], 'ndim':[1]}, 'dradius':{'max':lkok, 'min':['data'], 'ndim':[1,2]}, 'd0d':{'max':lkok, 'min':['data'], 'ndim':[1,2,3]}, @@ -2418,71 +2418,205 @@ def _checkformat_dtrm(dtime=None, dradius=None, dmesh=None, msg += " - Expected: %s\n"%str(dkok[dk]['ndim']) msg += " - Provided: %s"%str(dd[dk][k0]['data'].ndim) raise Exception(msg) + + # mesh if dk == 'dmesh': - dd[dk][k0]['nodes'] = np.atleast_2d(v0['nodes']).astype(float) - dd[dk][k0]['faces'] = np.atleast_2d(v0['faces']).astype(int) - nnodes = dd[dk][k0]['nodes'].shape[0] - nfaces = dd[dk][k0]['faces'].shape[0] - - # Test for duplicates - nodesu = np.unique(dd[dk][k0]['nodes'], axis=0) - facesu = np.unique(dd[dk][k0]['faces'], axis=0) - lc = [nodesu.shape[0] != nnodes, - facesu.shape[0] != nfaces] - if any(lc): - msg = "Non-valid mesh %s[%s]:\n"%(dk,k0) - if lc[0]: - msg += " Duplicate nodes: %s\n"%str(nnodes - nodesu.shape[0]) - msg += " - nodes.shape: %s\n"%str(dd[dk][k0]['nodes'].shape) - msg += " - unique nodes.shape: %s\n"%str(nodesu.shape) - if lc[1]: - msg += " Duplicate faces: %s\n"%str(nfaces - facesu.shape[0]) - msg += " - faces.shape: %s\n"%str(dd[dk][k0]['faces'].shape) - msg += " - unique faces.shape: %s"%str(facesu.shape) + + lmok = ['rect', 'tri', 'quadtri'] + if v0['type'] not in lmok: + msg = ("Mesh['type'] should be in {}\n".format(lmok) + + "\t- Provided: {}".format(v0['type'])) raise Exception(msg) - # Test for unused nodes - facesu = np.unique(facesu) - c0 = np.all(facesu>=0) and facesu.size == nnodes - if not c0: - indnot = [ii for ii in range(0,nnodes) - if ii not in facesu] - msg = "Some nodes not used in mesh %s[%s]:\n"(dk,k0) - msg += " - unused nodes indices: %s"%str(indnot) - warnings.warn(msg) - - - dd[dk][k0]['nnodes'] = dd[dk][k0].get('nnodes', nnodes) - dd[dk][k0]['nfaces'] = dd[dk][k0].get('nfaces', nfaces) - - assert dd[dk][k0]['nodes'].shape == (v0['nnodes'],2) - assert np.max(dd[dk][k0]['faces']) < v0['nnodes'] - # Only triangular meshes so far - assert v0['type'] in ['tri', 'quadtri'], v0['type'] - - if 'tri' in v0['type']: - assert dd[dk][k0]['faces'].shape == (v0['nfaces'],3) - if v0.get('mpltri', None) is None: - dd[dk][k0]['mpltri'] = mplTri(dd[dk][k0]['nodes'][:,0], - dd[dk][k0]['nodes'][:,1], - dd[dk][k0]['faces']) - assert isinstance(dd[dk][k0]['mpltri'], mplTri) - assert dd[dk][k0]['ftype'] in [0,1] - ntri = dd[dk][k0]['ntri'] - if dd[dk][k0]['ftype'] == 1: - dd[dk][k0]['size'] = dd[dk][k0]['nnodes'] + if v0['type'] == 'rect': + c0 = all([ss in v0.keys() and v0[ss].ndim in [1, 2] + for ss in ['R', 'Z']]) + if not c0: + msg = ("A mesh of type 'rect' must have attr.:\n" + + "\t- R of dim in [1, 2]\n" + + "\t- Z of dim in [1, 2]") + raise Exception(msg) + shapeu = np.unique(np.r_[v0['R'].shape, v0['Z'].shape]) + + shapeRZ = v0['shapeRZ'] + if shapeRZ is None: + shapeRZ = [None, None] else: - dd[dk][k0]['size'] = int(dd[dk][k0]['nfaces']/ntri) + shapeRZ = list(shapeRZ) + if v0['R'].ndim == 1: + if np.any(np.diff(v0['R']) <= 0.): + msg = "Non-increasing R" + raise Exception(msg) + R = v0['R'] + else: + lc = [np.all(np.diff(v0['R'][0, :])) > 0., + np.all(np.diff(v0['R'][:, 0])) > 0.] + if np.sum(lc) != 1: + msg = "Impossible to know R dimension!" + raise Exception(msg) + if lc[0]: + R = v0['R'][0, :] + if shapeRZ[1] is None: + shapeRZ[1] = 'R' + if shapeRZ[1] != 'R': + msg = "Inconsistent shapeRZ" + raise Exception(msg) + else: + R = v0['R'][:, 0] + if shapeRZ[0] is None: + shapeRZ[0] = 'R' + if shapeRZ[0] != 'R': + msg = "Inconsistent shapeRZ" + raise Exception(msg) + if v0['Z'].ndim == 1: + if np.any(np.diff(v0['Z']) <= 0.): + msg = "Non-increasing Z" + raise Exception(msg) + Z = v0['Z'] + else: + lc = [np.all(np.diff(v0['Z'][0, :])) > 0., + np.all(np.diff(v0['Z'][:, 0])) > 0.] + if np.sum(lc) != 1: + msg = "Impossible to know R dimension!" + raise Exception(msg) + if lc[0]: + Z = v0['Z'][0, :] + if shapeRZ[1] is None: + shapeRZ[1] = 'Z' + if shapeRZ[1] != 'Z': + msg = "Inconsistent shapeRZ" + raise Exception(msg) + else: + Z = v0['Z'][:, 0] + if shapeRZ[0] is None: + shapeRZ[0] = 'Z' + if shapeRZ[0] != 'Z': + msg = "Inconsistent shapeRZ" + raise Exception(msg) + shapeRZ = tuple(shapeRZ) + if shapeRZ not in [('R', 'Z'), ('Z', 'R')]: + msg = "Inconsistent shapeRZ" + raise Exception(msg) + + if None in shapeRZ: + msg = ("Please provide shapeRZ " + + " = ('R', 'Z') or ('Z', 'R')\n" + + "Could not be inferred from data itself") + raise Exception(msg) + + def trifind(r, z, + Rbin=0.5*(R[1:] + R[:-1]), + Zbin=0.5*(Z[1:] + Z[:-1]), + nR=R.size, nZ=Z.size, + shapeRZ=shapeRZ): + indR = np.searchsorted(Rbin, r) + indZ = np.searchsorted(Zbin, z) + if shapeRZ == ('R', 'Z'): + indpts = indR*nZ + indZ + else: + indpts = indZ*nR + indR + return indpts + + dd[dk][k0]['R'] = R + dd[dk][k0]['Z'] = Z + dd[dk][k0]['shapeRZ'] = shapeRZ + dd[dk][k0]['nR'] = R.size + dd[dk][k0]['nZ'] = Z.size + dd[dk][k0]['trifind'] = trifind + if dd[dk][k0]['ftype'] != 0: + msg = "Linear interpolation not handled yet !" + raise Exception(msg) + dd[dk][k0]['size'] = R.size*Z.size + + else: + ls = ['nodes', 'faces'] + if not all([s in v0.keys() for s in ls]): + msg = ("The following keys should be in dmesh:\n" + + "\t- {}".format(ls)) + raise Exception(msg) + func = np.atleast_2d + dd[dk][k0]['nodes'] = func(v0['nodes']).astype(float) + dd[dk][k0]['faces'] = func(v0['faces']).astype(int) + nnodes = dd[dk][k0]['nodes'].shape[0] + nfaces = dd[dk][k0]['faces'].shape[0] + + # Test for duplicates + nodesu = np.unique(dd[dk][k0]['nodes'], axis=0) + facesu = np.unique(dd[dk][k0]['faces'], axis=0) + lc = [nodesu.shape[0] != nnodes, + facesu.shape[0] != nfaces] + if any(lc): + msg = "Non-valid mesh {0}[{1}]: \n".format(dk, k0) + if lc[0]: + ndup = nnodes - nodesu.shape[0] + ndsh = dd[dk][k0]['nodes'].shape + undsh = nodesu.shape + msg += ( + " Duplicate nodes: {}\n".format(ndup) + + "\t- nodes.shape: {}\n".format(nodsh) + + "\t- unique shape: {}\n".format(undsh)) + if lc[1]: + ndup = str(nfaces - facesu.shape[0]) + facsh = str(dd[dk][k0]['faces'].shape) + ufacsh = str(facesu.shape) + msg += ( + " Duplicate faces: {}\n".format(ndup) + + "\t- faces.shape: {}\n".format(facsh) + + "\t- unique shape: {}".format(ufacsh)) + raise Exception(msg) + + # Test for unused nodes + facesu = np.unique(facesu) + c0 = np.all(facesu >= 0) and facesu.size == nnodes + if not c0: + ino = str([ii for ii in range(0, nnodes) + if ii not in facesu]) + msg = "Unused nodes in {0}[{1}]:\n".format(dk, k0) + msg += " - unused nodes indices: {}".format(ino) + warnings.warn(msg) + + dd[dk][k0]['nnodes'] = dd[dk][k0].get('nnodes', nnodes) + dd[dk][k0]['nfaces'] = dd[dk][k0].get('nfaces', nfaces) + + assert dd[dk][k0]['nodes'].shape == (v0['nnodes'], 2) + assert np.max(dd[dk][k0]['faces']) < v0['nnodes'] + # Only triangular meshes so far + assert v0['type'] in ['tri', 'quadtri'], v0['type'] + + if 'tri' in v0['type']: + fshap = dd[dk][k0]['faces'].shape + fshap0 = (v0['nfaces'], 3) + if fshap != fshap0: + msg = ("Wrong shape of {}[{}]\n".format(dk, k0) + + "\t- Expected: {}\n".format(fshap0) + + "\t- Provided: {}".format(fshap)) + raise Exception(msg) + if v0.get('mpltri', None) is None: + dd[dk][k0]['mpltri'] = mplTri( + dd[dk][k0]['nodes'][:, 0], + dd[dk][k0]['nodes'][:, 1], + dd[dk][k0]['faces']) + assert isinstance(dd[dk][k0]['mpltri'], mplTri) + assert dd[dk][k0]['ftype'] in [0, 1] + ntri = dd[dk][k0]['ntri'] + if dd[dk][k0]['ftype'] == 1: + dd[dk][k0]['size'] = dd[dk][k0]['nnodes'] + else: + dd[dk][k0]['size'] = int( + dd[dk][k0]['nfaces'] / ntri + ) # Check unicity of all keys lk = [list(dv.keys()) for dv in dd.values()] lk = list(itt.chain.from_iterable(lk)) lku = sorted(set(lk)) - lk = ['%s : %s times'%(kk, str(lk.count(kk))) for kk in lku if lk.count(kk) > 1] + lk = ['{0} : {1} times'.format(kk, str(lk.count(kk))) + for kk in lku if lk.count(kk) > 1] if len(lk) > 0: - msg = "Each key of (dtime,dradius,dmesh,d0d,d1d,d2d) must be unique !\n" - msg += "The following keys are repeated :\n" - msg += " - " + "\n - ".join(lk) + msg = ("Each key of (dtime, dradius, dmesh, d0d, d1d, d2d)" + + " must be unique !\n" + + "The following keys are repeated :\n" + + " - " + "\n - ".join(lk)) raise Exception(msg) dtime, dradius, dmesh = dd['dtime'], dd['dradius'], dd['dmesh'] @@ -2634,6 +2768,10 @@ def _set_dindrefdatagroup(self, dtime=None, dradius=None, dmesh=None, lkt = [k for k in dtime.keys() if k in dradius[k0]['depend']] assert len(lkt) == 1 axist = dradius[k0]['depend'].index(lkt[0]) + # Handle cases with only 1 time step + if data.ndim == 1: + assert dindref[lkt[0]]['size'] == 1 + data = data.reshape((1, data.size)) size = data.shape[1-axist] dindref[k0] = {'size':size, 'group':'radius'} @@ -3512,13 +3650,22 @@ def _get_indtmult(self, idquant=None, idref1d=None, idref2d=None): # Get indtqr1r2 (tall with respect to tq, tr1, tr2) indtq, indtr1, indtr2 = None, None, None - indtq = np.digitize(tall, tbinq) + if tbinq.size > 0: + indtq = np.digitize(tall, tbinq) + else: + indtq = np.r_[0] if idref1d is None: assert np.all(indtq == np.arange(0,tall.size)) if idref1d is not None: - indtr1 = np.digitize(tall, tbinr1) + if tbinr1.size > 0: + indtr1 = np.digitize(tall, tbinr1) + else: + indtr1 = np.r_[0] if idref2d is not None: - indtr2 = np.digitize(tall, tbinr2) + if tbinr2.size > 0: + indtr2 = np.digitize(tall, tbinr2) + else: + indtr2 = np.r_[0] ntall = tall.size return tall, tbinall, ntall, indtq, indtr1, indtr2 @@ -3604,14 +3751,18 @@ def _get_finterp(self, assert len(lidmesh) == 1 idmesh = lidmesh[0] - # Get mesh - mpltri = self._ddata[idmesh]['data']['mpltri'] - trifind = mpltri.get_trifinder() - # Get common time indices if interp_t == 'nearest': - out = self._get_tcom(idquant,idref1d, idref2d, idq2dR) - tall, tbinall, ntall, indtq, indtr1, indtr2= out + out = self._get_tcom(idquant, idref1d, idref2d, idq2dR) + tall, tbinall, ntall, indtq, indtr1, indtr2 = out + + # Get mesh + if self._ddata[idmesh]['data']['type'] == 'rect': + mpltri = None + trifind = self._ddata[idmesh]['data']['trifind'] + else: + mpltri = self._ddata[idmesh]['data']['mpltri'] + trifind = mpltri.get_trifinder() # # Prepare output @@ -3619,7 +3770,9 @@ def _get_finterp(self, # Note : Maybe consider using scipy.LinearNDInterpolator ? if idquant is not None: vquant = self._ddata[idquant]['data'] - if self._ddata[idmesh]['data']['ntri'] > 1: + c0 = (self._ddata[idmesh]['data']['type'] == 'quadtri' + and self._ddata[idmesh]['data']['ntri'] > 1) + if c0: vquant = np.repeat(vquant, self._ddata[idmesh]['data']['ntri'], axis=0) else: @@ -3654,7 +3807,6 @@ def _get_finterp(self, indtq=indtq, indtr1=indtr1, indtr2=indtr2, trifind=trifind) - return func diff --git a/tofu/geom/_core.py b/tofu/geom/_core.py index 60ce6d4bd..123c9fd49 100644 --- a/tofu/geom/_core.py +++ b/tofu/geom/_core.py @@ -5241,17 +5241,21 @@ def get_subset(self, indch=None, Name=None): if indch is None: return self else: + indch = self._check_indch(indch) dd = self.to_dict() + sep = [kk for kk in dd.keys() + if all([ss in kk for ss in ['dId', 'dall', 'Name']])][0] + sep = sep[3] # Name assert Name in [None, True] or type(Name) is str if Name is True: pass elif type(Name) is str: - dd["dId_dall_Name"] = Name + dd[sep.join(['dId', 'dall', 'Name'])] = Name elif Name is None: - dd["dId_dall_Name"] = dd["dId_dall_Name"] + "-subset" + dd[sep.join(['dId', 'dall', 'Name'])] += "-subset" # Resize all np.ndarrays for kk in dd.keys(): @@ -5262,7 +5266,8 @@ def get_subset(self, indch=None, Name=None): dd[kk] = vv[indch] elif vv.ndim == 2 and vv.shape[1] == self.nRays: dd[kk] = vv[:, indch] - dd["dgeom_nRays"] = dd["dgeom_D"].shape[1] + dd[sep.join(['dgeom', 'nRays'])] = ( + dd[sep.join(['dgeom', 'D'])].shape[1]) # Recreate from dict obj = self.__class__(fromdict=dd) diff --git a/tofu/imas2tofu/_core.py b/tofu/imas2tofu/_core.py index da640dc47..39ee55ce1 100644 --- a/tofu/imas2tofu/_core.py +++ b/tofu/imas2tofu/_core.py @@ -1804,8 +1804,9 @@ def _set_fsig(self): for k,v in self._dshort[ids].items(): self._dshort[ids][k]['fsig'] = self._get_fsig(v['str']) - def _get_data(self, ids, sig, occ, comp=False, indt=None, indch=None, - stack=True, isclose=True, flatocc=True, nan=True, pos=None): + def __get_data(self, ids, sig, occ, comp=False, indt=None, indch=None, + stack=True, isclose=True, flatocc=True, nan=True, + pos=None, warn=True): # get list of results for occ occref = self._dids[ids]['occ'] @@ -1814,9 +1815,10 @@ def _get_data(self, ids, sig, occ, comp=False, indt=None, indch=None, if comp: lstr = self._dcomp[ids][sig]['lstr'] kargs = self._dcomp[ids][sig].get('kargs', {}) - ddata = self.get_data(ids=ids, sig=lstr, - occ=occ, indch=indch, indt=indt, - stack=stack, flatocc=False, nan=nan, pos=pos) + ddata, _ = self._get_data(ids=ids, sig=lstr, + occ=occ, indch=indch, indt=indt, + stack=stack, flatocc=False, nan=nan, + pos=pos, warn=warn) out = [self._dcomp[ids][sig]['func']( *[ddata[kk][nn] for kk in lstr], **kargs ) for nn in range(0,nocc)] @@ -1853,15 +1855,9 @@ def _get_data(self, ids, sig, occ, comp=False, indt=None, indch=None, out = out[0] return out - def get_data(self, ids=None, sig=None, occ=None, - indch=None, indt=None, stack=True, - isclose=None, flatocc=True, nan=True, pos=None): - """ Return a dict of the desired signals extracted from specified ids - - If the ids has a field 'channel', indch is used to specify from which - channel data shall be loaded (all by default) - - """ + def _get_data(self, ids=None, sig=None, occ=None, + indch=None, indt=None, stack=True, + isclose=None, flatocc=True, nan=True, pos=None, warn=True): # ------------------ # Check format input @@ -1894,24 +1890,39 @@ def get_data(self, ids=None, sig=None, occ=None, # ------------------ # get data - dout = dict.fromkeys(sig) - for ii in range(0,len(sig)): + dout, dfail = {}, {} + for ii in range(0, len(sig)): if isclose is None: isclose_ = sig[ii] == 't' else: isclose_ = isclose try: - dout[sig[ii]] = self._get_data(ids, sig[ii], occ, comp=comp[ii], - indt=indt, indch=indch, - stack=stack, isclose=isclose_, - flatocc=flatocc, nan=nan, - pos=pos) + dout[sig[ii]] = self.__get_data(ids, sig[ii], occ, + comp=comp[ii], + indt=indt, indch=indch, + stack=stack, isclose=isclose_, + flatocc=flatocc, nan=nan, + pos=pos, warn=warn) except Exception as err: - msg = '\n' + str(err) + '\n' - msg += '\tIn ids %s, signal %s not loaded !'%(ids,sig[ii]) - warnings.warn(msg) - del dout[sig[ii]] - return dout + dfail[sig[ii]] = str(err) + if warn: + msg = '\n' + str(err) + '\n' + msg += '\tsignal {0}.{1} not loaded!'.format(ids, sig[ii]) + warnings.warn(msg) + return dout, dfail + + def get_data(self, ids=None, sig=None, occ=None, + indch=None, indt=None, stack=True, + isclose=None, flatocc=True, nan=True, pos=None, warn=True): + """ Return a dict of the desired signals extracted from specified ids + + If the ids has a field 'channel', indch is used to specify from which + channel data shall be loaded (all by default) + + """ + return self._get_data(ids=ids, sig=sig, occ=occ, indch=indch, + indt=indt, stack=stack, isclose=isclose, + flatocc=flatocc, nan=nan, pos=pos, warn=warn)[0] def get_data_all(self, dsig=None, stack=True, isclose=None, flatocc=True, nan=True, pos=None): @@ -1934,14 +1945,31 @@ def get_data_all(self, dsig=None, stack=True, assert all([type(v) in [str,list] or v is None for v in dsig.values()]) # Get data + dfail = dict.fromkeys(dout.keys()) + anyerror = False for ids in dout.keys(): try: - dout[ids] = self.get_data(ids, sig=dsig[ids], stack=stack, - isclose=isclose, flatocc=flatocc, - nan=nan, pos=pos) + dout[ids], dfail[ids] = self._get_data(ids, sig=dsig[ids], + stack=stack, + isclose=isclose, + flatocc=flatocc, + nan=nan, + pos=pos, warn=False) + if len(dfail[ids]) > 0: + anyerror = True except Exception as err: - msg = "Could not get data from %s"%ids - warnings.warn(msg) + del dout[ids] + dfail[ids] = dict.fromkeys(dsig[ids].keys(), 'ids error') + anyerror = True + if anyerror: + msg = "The following data could not be retrieved:" + for ids, v0 in dfail.items(): + if len(v0) == 0: + continue + msg += "\n\t- {}:".format(ids) + for k1, v1 in v0.items(): + msg += "\n\t\t{0} : {1}".format(k1, v1.replace('\n', ' ')) + warnings.warn(msg) return dout @@ -2179,13 +2207,13 @@ def _checkformat_Plasma2D_dsig(self, dsig=None): @staticmethod - def _checkformat_mesh(nodes, indfaces, ids=None): + def _checkformat_mesh_NodesFaces(nodes, indfaces, ids=None): # Check mesh type if indfaces.shape[1] == 3: - meshtype = 'tri' + mtype = 'tri' elif indfaces.shape[1] == 4: - meshtype = 'quad' + mtype = 'quad' else: msg = "Mesh seems to be neither triangular nor quadrilateral\n" msg += " => unrecognized mesh type, not implemented yet" @@ -2246,7 +2274,7 @@ def _checkformat_mesh(nodes, indfaces, ids=None): warnings.warn(msg) # Convert to triangular mesh if necessary - if meshtype == 'quad': + if mtype == 'quad': # Convert to tri mesh (solution for unstructured meshes) indface = np.empty((indfaces.shape[0]*2,3), dtype=int) @@ -2254,7 +2282,7 @@ def _checkformat_mesh(nodes, indfaces, ids=None): indface[1::2,:-1] = indfaces[:,2:] indface[1::2,-1] = indfaces[:,0] indfaces = indface - meshtype = 'quadtri' + mtype = 'quadtri' ntri = 2 else: ntri = 1 @@ -2274,11 +2302,10 @@ def _checkformat_mesh(nodes, indfaces, ids=None): warnings.warn(msg) (indfaces[indclock,1], indfaces[indclock,2]) = indfaces[indclock,2], indfaces[indclock,1] - return indfaces, meshtype, ntri - + return indfaces, mtype, ntri - # TBF - def inspect_ggd(ids): + def inspect_ggd(self, ids): + # TBF if ids not in self._dids.keys(): msg = "The ggd of ids %s cannot be inspected:\n"%ids msg += " => please add ids first (self.add_ids())" @@ -2300,6 +2327,61 @@ def inspect_ggd(ids): for ll in range(0,nspace): npts = ggd.space[ll].objects_per_dimension[0].object[0] + @staticmethod + def _checkformat_mesh_Rect(R, Z, datashape=None, + shapeRZ=None, ids=None): + assert R.ndim in [1, 2] + assert Z.ndim in [1, 2] + shapeu = np.unique(np.r_[R.shape, Z.shape]) + shapeRZ = [None, None] + if R.ndim == 1: + assert np.all(np.diff(R) > 0.) + else: + lc = [np.all(np.diff(R[0, :])) > 0., + np.all(np.diff(R[:, 0])) > 0.] + assert np.sum(lc) == 1 + if lc[0]: + R = R[0, :] + if shapeRZ[1] is None: + shapeRZ[1] = 'R' + assert shapeRZ[1] == 'R' + else: + R = R[:, 0] + if shapeRZ[0] is None: + shapeRZ[0] = 'R' + assert shapeRZ[0] == 'R' + if Z.ndim == 1: + assert np.all(np.diff(Z) > 0.) + else: + lc = [np.all(np.diff(Z[0, :])) > 0., + np.all(np.diff(Z[:, 0])) > 0.] + assert np.sum(lc) == 1 + if lc[0]: + Z = Z[0, :] + if shapeRZ[1] is None: + shapeRZ[1] = 'Z' + assert shapeRZ[1] == 'Z' + else: + Z = Z[:, 0] + if shapeRZ[0] is None: + shapeRZ[0] = 'Z' + assert shapeRZ[0] == 'Z' + + if datashape is not None: + if None in shapeRZ: + pass + + if shapeRZ == ['R', 'Z']: + assert datashape == (R.size, Z.size) + elif shapeRZ == ['Z', 'R']: + assert datashape == (Z.size, R.size) + else: + msg = "Inconsistent data shape !" + raise Exception(msg) + + shapeRZ = tuple(shapeRZ) + assert shapeRZ in [('R', 'Z'), ('Z', 'R')] + return R, Z, shapeRZ, 0 @@ -2441,7 +2523,7 @@ def _get_dextra(self, dextra=None, fordata=False, nan=True, pos=None): def to_Plasma2D(self, tlim=None, dsig=None, t0=None, Name=None, occ=None, config=None, out=object, plot=None, plot_sig=None, plot_X=None, - bck=True, dextra=None, nan=True, pos=None): + bck=True, dextra=None, nan=True, pos=None, shapeRZ=None): # dsig dsig = self._checkformat_Plasma2D_dsig(dsig) @@ -2490,7 +2572,7 @@ def to_Plasma2D(self, tlim=None, dsig=None, t0=None, _, _, shot, Exp = self._get_lidsidd_shotExp(lids, upper=True, errshot=True, errExp=True) # get data - out_ = self.get_data_all(dsig=dsig) + out0 = self.get_data_all(dsig=dsig) # ------------- # Input dicts @@ -2507,10 +2589,14 @@ def to_Plasma2D(self, tlim=None, dsig=None, t0=None, d1d, dradius = {}, {} d2d, dmesh = {}, {} for ids in lids: + # Hotfiw to avoid calling get_data a second time out0 -> out_ + # TBF in next release (ugly, sub-optimal...) # dtime - out_ = self.get_data(ids, sig='t') - lc = [len(out_) == 1, out_['t'].size > 0, 0 not in out_['t'].shape] + out_ = {'t': out0[ids].get('t', None)} + lc = [out_['t'] is not None, + out_['t'].size > 0, + 0 not in out_['t'].shape] keyt, nt, indt = None, None, None if all(lc): nt = out_['t'].size @@ -2520,10 +2606,14 @@ def to_Plasma2D(self, tlim=None, dsig=None, t0=None, dtime[keyt] = {'data':dtt['t'], 'origin':ids, 'name':'t'} indt = dtt['indt'] + else: + nt = None # d1d and dradius - lsig = [k for k in dsig[ids] if '1d' in k] - out_ = self.get_data(ids, lsig, indt=indt, nan=nan, pos=pos) + lsig = [k for k in dsig[ids] + if '1d' in k and k in out0[ids].keys()] + out_ = self.get_data(ids, lsig, indt=indt, nan=nan, pos=pos, + warn=False) if len(out_) > 0: nref, kref = None, None for ss in out_.keys(): @@ -2549,10 +2639,10 @@ def to_Plasma2D(self, tlim=None, dsig=None, t0=None, msg += " - %s.%s.shape = %s"%(ids,ss,str(shape)) msg += " - One dim should be t.size = %s"%str(nt) raise Exception(msg) - axist = shape.index(nt) - nr = shape[1-axist] - if axist == 1: - out_[ss] = out_[ss].T + axist = shape.index(nt) + nr = shape[1-axist] + if axist == 1: + out_[ss] = out_[ss].T if ss in self._dshort[ids].keys(): dim = self._dshort[ids][ss].get('dim', 'unknown') @@ -2565,16 +2655,17 @@ def to_Plasma2D(self, tlim=None, dsig=None, t0=None, key = '%s.%s'%(ids,ss) if nref is None: - dradius[key] = {'data':out_[ss], 'name':ss, - 'origin':ids, 'dim':dim, 'quant':quant, - 'units':units, 'depend':(keyt,key)} + dradius[key] = {'data': out_[ss], 'name': ss, + 'origin': ids, 'dim': dim, + 'quant': quant, 'units': units, + 'depend': (keyt, key)} nref, kref = nr, key else: assert nr == nref - d1d[key] = {'data':out_[ss], 'name':ss, - 'origin':ids, 'dim':dim, 'quant':quant, - 'units':units, 'depend':(keyt,kref)} - assert out_[ss].shape == (nt,nr) + d1d[key] = {'data': out_[ss], 'name': ss, + 'origin': ids, 'dim': dim, 'quant': quant, + 'units': units, 'depend': (keyt, kref)} + assert out_[ss].shape == (nt, nr) if plot: if ss in plot_sig: @@ -2583,22 +2674,53 @@ def to_Plasma2D(self, tlim=None, dsig=None, t0=None, plot_X[plot_X.index(ss)] = key # d2d and dmesh - lsig = [k for k in dsig[ids] if '2d' in k] - lsigmesh = ['2dmeshNodes','2dmeshFaces'] - out_ = self.get_data(ids, sig=lsig, indt=indt, nan=nan, pos=pos) + lsig = [k for k in dsig[ids] + if '2d' in k and k in out0[ids].keys()] + lsigmesh = [k for k in lsig if 'mesh' in k] + out_ = self.get_data(ids, sig=lsig, indt=indt, nan=nan, pos=pos, + warn=False) - lc = [len(out_) > 0, all([ss in out_.keys() for ss in lsigmesh])] - if all(lc): - npts = None - keym = '%s.mesh'%ids + cmesh = any([ss in out_.keys() for ss in lsigmesh]) + if len(out_) > 0: + npts, datashape = None, None + keym = '{}.mesh'.format(ids) if cmesh else None for ss in set(out_.keys()).difference(lsigmesh): assert out_[ss].ndim in [1,2] if out_[ss].ndim == 1: out_[ss] = np.atleast_2d(out_[ss]) shape = out_[ss].shape - assert len(shape) == 2 + assert len(shape) >= 2 if np.sum(shape) > 0: - assert nt in shape + if nt not in shape: + assert nt == 1 + assert len(shape) == 2 + datashape = tuple(shape) + if shapeRZ is None: + msg = ("Please provide shapeRZ" + + "indexing is ambiguous") + raise Exception(msg) + size = shape[0]*shape[1] + if shapeRZ == ('R', 'Z'): + out_[ss] = out_[ss].reshape((nt, size)) + elif shapeRZ == ('Z', 'R'): + out_[ss] = out_[ss].reshape((nt, size), + order='F') + shape = out_[ss].shape + if len(shape) == 3: + assert nt == shape[0] + datashape = tuple(shape[1], shape[2]) + if shapeRZ is None: + msg = ("Please provide shapeRZ" + + "indexing is ambiguous") + raise Exception(msg) + size = shape[1]*shape[2] + if shapeRZ == ('R', 'Z'): + out_[ss] = out_[ss].reshape((nt, size)) + elif shapeRZ == ('Z', 'R'): + out_[ss] = out_[ss].reshape((nt, size), + order='F') + shape = out_[ss].shape + axist = shape.index(nt) if npts is None: npts = shape[1-axist] @@ -2616,31 +2738,59 @@ def to_Plasma2D(self, tlim=None, dsig=None, t0=None, units = self._dcomp[ids][ss].get('units', 'a.u.') key = '%s.%s'%(ids,ss) - d2d[key] = {'data':out_[ss], 'name':ss, - 'dim':dim, 'quant':quant, 'units':units, - 'origin':ids, 'depend':(keyt,keym)} - - nodes = out_['2dmeshNodes'] - indfaces = out_['2dmeshFaces'] - indfaces,meshtype,ntri = self._checkformat_mesh(nodes,indfaces, - ids=ids) - nnod, nfaces = int(nodes.size/2), indfaces.shape[0] - if npts is not None: - if npts not in [nnod, int(nfaces/ntri)]: - msg = "There is an indexing unconsistency:\n" - msg += " - 2d profiles have npts = %s\n"%str(npts) - msg += " - mesh has %s nodes\n"%str(nnod) - msg += " %s faces"%str(int(nfaces/ntri)) + d2d[key] = {'data': out_[ss], 'name': ss, + 'dim': dim, 'quant': quant, 'units': units, + 'origin': ids, 'depend': (keyt, keym)} + + if cmesh: + lc = [all([ss in lsig for ss in ['2dmeshNodes', + '2dmeshFaces']]), + all([ss in lsig for ss in ['2dmeshR', '2dmeshZ']])] + if not np.sum(lc) == 1: + msg = ("2d mesh shall be provided either via:\n" + + "\t- '2dmeshR' and '2dmeshZ'\n" + + "\t- '2dmeshNodes' and '2dmeshFaces'") raise Exception(msg) - ftype = 1 if npts == nnod else 0 - else: - ftype = None - mpltri = mpl.tri.Triangulation(nodes[:,0], nodes[:,1], indfaces) - dmesh[keym] = {'dim':'mesh', 'quant':'mesh', 'units':'a.u.', - 'origin':ids, 'depend':(keym,), 'name':meshtype, - 'nodes':nodes, 'faces':indfaces, - 'type':meshtype, 'ntri':ntri, 'ftype':ftype, - 'nnodes':nnod,'nfaces':nfaces, 'mpltri':mpltri} + + # Nodes / Faces case + if lc[0]: + nodes = out_['2dmeshNodes'] + indfaces = out_['2dmeshFaces'] + func = self._checkformat_mesh_NodesFaces + indfaces, mtype, ntri = func(nodes, indfaces, ids=ids) + nnod, nfaces = int(nodes.size/2), indfaces.shape[0] + if npts is not None: + nft = int(nfaces/ntri) + if npts not in [nnod, nft]: + msg = "Inconsistent indices:\n" + msg += "\t- 2d profile {} npts\n".format(npts) + msg += "\t- mesh {} nodes\n".format(nnod) + msg += "\t {} faces".format(nft) + raise Exception(msg) + ftype = 1 if npts == nnod else 0 + else: + ftype = None + mpltri = mpl.tri.Triangulation(nodes[:, 0], + nodes[:, 1], indfaces) + dmesh[keym] = {'dim': 'mesh', 'quant': 'mesh', + 'units': 'a.u.', 'origin': ids, + 'depend': (keym,), 'name': mtype, + 'nodes': nodes, 'faces': indfaces, + 'type': mtype, 'ntri': ntri, + 'ftype': ftype, 'nnodes': nnod, + 'nfaces': nfaces, 'mpltri': mpltri} + # R / Z case + elif lc[1]: + func = self._checkformat_mesh_Rect + R, Z, shapeRZ, ftype = func(out_['2dmeshR'], + out_['2dmeshZ'], + datashape=datashape, + shapeRZ=shapeRZ, ids=ids) + dmesh[keym] = {'dim': 'mesh', 'quant': 'mesh', + 'units': 'a.u.', 'origin': ids, + 'depend': (keym,), 'name': 'rect', + 'R': R, 'Z': Z, 'shapeRZ': shapeRZ, + 'type': 'rect', 'ftype': ftype} # t0 t0 = self._get_t0(t0) @@ -2648,7 +2798,6 @@ def to_Plasma2D(self, tlim=None, dsig=None, t0=None, for tt in dtime.keys(): dtime[tt]['data'] = dtime[tt]['data'] - t0 - plasma = dict(dtime=dtime, dradius=dradius, dmesh=dmesh, d0d=d0d, d1d=d1d, d2d=d2d, Exp=Exp, shot=shot, Name=Name, config=config) @@ -2988,14 +3137,12 @@ def to_Data(self, ids=None, dsig=None, data=None, X=None, tlim=None, msg += "\n => Solution: choose indch accordingly !" raise Exception(msg) - if t.ndim == 2: assert np.all(np.isclose(t, t[0:1,:])) t = t[0,:] dins['t'] = t indt = self._checkformat_tlim(t, tlim=tlim)['indt'] - out = self.get_data(ids, sig=[dsig[k] for k in lk], indt=indt, indch=indch, nan=nan, pos=pos) for kk in set(lk).difference('t'): @@ -3071,7 +3218,8 @@ def to_Data(self, ids=None, dsig=None, data=None, X=None, tlim=None, if 'validity_timed' in self._dshort[ids].keys(): inan = self.get_data(ids, sig='validity_timed', - indt=indt, indch=indch, nan=nan, pos=pos)['validity_timed'].T<0. + indt=indt, indch=indch, + nan=nan, pos=pos)['validity_timed'].T < 0. dins['data'][inan] = np.nan if 'X' in dins.keys() and np.any(np.isnan(dins['X'])): if fallback_X is None: diff --git a/tofu/version.py b/tofu/version.py index d300cef4c..38f81c4e7 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-alpha2-28-gb8322fe' +__version__ = '1.4.2-alpha2-48-gbe9fde3'