diff --git a/fluents/dataset.py b/fluents/dataset.py index d871b4e..61af574 100644 --- a/fluents/dataset.py +++ b/fluents/dataset.py @@ -224,6 +224,22 @@ class Dataset(object): """ return copy.deepcopy(self) + def subdata(self, dim, idents): + """Returns a new dataset based on dimension and given identifiers. + + """ + ds = self.copy() + indices = ds.get_indices(dim, idents) + ax = [i for i, name in enumerate(ds._dims) if name == dim][0] + subarr = ds._array.take(indices, ax) + for k, v in ds._map[dim].items(): + if k not in idents: + del ds._map[dim][k] + ds.shape = tuple(len(ds._map[d]) for d in ds._dims) + ds.set_array(subarr) + return ds + + def transpose(self): """Returns a copy of transpose of a dataset. @@ -569,7 +585,7 @@ def write_ftsv(fd, ds, decimals=7, sep='\t', fmt=None, sp_format=True): fd.write("# name: %s" % ds.get_name() + '\n') # xy-node-positions - if ds.nodepos != None: + if type == 'network' and ds.nodepos != None: fd.write("# nodepos:") node_dim = ds.get_dim_name(0) for ident in ds.get_identifiers(node_dim, sorted=True):