|
@ -36,6 +36,17 @@ class MongoDB: |
|
|
except: |
|
|
except: |
|
|
pass |
|
|
pass |
|
|
return data |
|
|
return data |
|
|
|
|
|
|
|
|
|
|
|
def _break_dataSet(self, dataSet, scanAxis=None): |
|
|
|
|
|
|
|
|
|
|
|
if scanAxis is None: |
|
|
|
|
|
scanAxis = dataSet.attrs['scanAxis'] |
|
|
|
|
|
|
|
|
|
|
|
dataArray = dataSet.shotNum |
|
|
|
|
|
|
|
|
|
|
|
stackedDataArray = dataArray.stack(_scanAxis=tuple(scanAxis)) |
|
|
|
|
|
|
|
|
|
|
|
return stackedDataArray |
|
|
|
|
|
|
|
|
def set_date(self, date): |
|
|
def set_date(self, date): |
|
|
date = date.split("/") |
|
|
date = date.split("/") |
|
@ -52,6 +63,7 @@ class MongoDB: |
|
|
'month': self.month, |
|
|
'month': self.month, |
|
|
'day': self.day, |
|
|
'day': self.day, |
|
|
'shotNum': shotNum, |
|
|
'shotNum': shotNum, |
|
|
|
|
|
'runNum': 0, |
|
|
'global_parameters' : {}, |
|
|
'global_parameters' : {}, |
|
|
} |
|
|
} |
|
|
|
|
|
|
|
@ -62,43 +74,117 @@ class MongoDB: |
|
|
|
|
|
|
|
|
data = self._convert_numpy_type(data) |
|
|
data = self._convert_numpy_type(data) |
|
|
|
|
|
|
|
|
self.mongoDB['global'].insert_one(data) |
|
|
|
|
|
|
|
|
if 'scanAxis' in dataSet.attrs: |
|
|
|
|
|
|
|
|
|
|
|
del data['global_parameters']['scanAxis'] |
|
|
|
|
|
del data['global_parameters']['scanAxisLength'] |
|
|
|
|
|
|
|
|
|
|
|
scanAxis = dataSet.attrs['scanAxis'] |
|
|
|
|
|
data['global_parameters'].update( |
|
|
|
|
|
{ |
|
|
|
|
|
key:0 |
|
|
|
|
|
for key in scanAxis |
|
|
|
|
|
} |
|
|
|
|
|
) |
|
|
|
|
|
|
|
|
|
|
|
stackedDataArray = self._break_dataSet(dataSet) |
|
|
|
|
|
stackedDataArray = stackedDataArray.groupby('_scanAxis') |
|
|
|
|
|
|
|
|
|
|
|
for i in stackedDataArray: |
|
|
|
|
|
|
|
|
|
|
|
stackedDataArray_single = i[1] |
|
|
|
|
|
|
|
|
|
|
|
data.update( |
|
|
|
|
|
{ |
|
|
|
|
|
'runNum': int(stackedDataArray_single.item()) |
|
|
|
|
|
} |
|
|
|
|
|
) |
|
|
|
|
|
|
|
|
|
|
|
data['global_parameters'].update( |
|
|
|
|
|
{ |
|
|
|
|
|
key: stackedDataArray_single[key].item() |
|
|
|
|
|
for key in scanAxis |
|
|
|
|
|
} |
|
|
|
|
|
) |
|
|
|
|
|
|
|
|
|
|
|
if '_id' in data: |
|
|
|
|
|
del data['_id'] |
|
|
|
|
|
|
|
|
|
|
|
self.mongoDB['global'].insert_one(data) |
|
|
|
|
|
|
|
|
|
|
|
else: |
|
|
|
|
|
self.mongoDB['global'].insert_one(data) |
|
|
|
|
|
|
|
|
def _add_data_normal(self, shotNum, data): |
|
|
|
|
|
|
|
|
def _add_data_normal(self, shotNum, runNum, data): |
|
|
|
|
|
|
|
|
|
|
|
if runNum is None: |
|
|
|
|
|
runNum = 0 |
|
|
|
|
|
|
|
|
filter = { |
|
|
filter = { |
|
|
'year': self.year, |
|
|
'year': self.year, |
|
|
'month': self.month, |
|
|
'month': self.month, |
|
|
'day': self.day, |
|
|
'day': self.day, |
|
|
'shotNum': shotNum, |
|
|
'shotNum': shotNum, |
|
|
|
|
|
'runNum': runNum, |
|
|
} |
|
|
} |
|
|
|
|
|
|
|
|
self.mongoDB['global'].update_one(filter, {"$set": data}, upsert=False) |
|
|
self.mongoDB['global'].update_one(filter, {"$set": data}, upsert=False) |
|
|
|
|
|
|
|
|
def _add_data_xarray_dataArray(self, shotNum, dataArray): |
|
|
|
|
|
filter = { |
|
|
|
|
|
'year': self.year, |
|
|
|
|
|
'month': self.month, |
|
|
|
|
|
'day': self.day, |
|
|
|
|
|
'shotNum': shotNum, |
|
|
|
|
|
} |
|
|
|
|
|
|
|
|
def _add_data_xarray_dataArray(self, shotNum, dataArray, scanAxis=None): |
|
|
|
|
|
|
|
|
|
|
|
if scanAxis is None: |
|
|
|
|
|
scanAxis = list(dataArray.coords) |
|
|
|
|
|
|
|
|
dataArray.attrs = self._convert_numpy_type(dataArray.attrs) |
|
|
dataArray.attrs = self._convert_numpy_type(dataArray.attrs) |
|
|
|
|
|
|
|
|
|
|
|
stackedDataArray = dataArray.stack(_scanAxis=tuple(scanAxis)) |
|
|
|
|
|
stackedDataArray = stackedDataArray.groupby('_scanAxis') |
|
|
|
|
|
|
|
|
mongoID, _ = self.xdb.put(dataArray) |
|
|
|
|
|
|
|
|
|
|
|
data_label = { |
|
|
|
|
|
dataArray.name: |
|
|
|
|
|
{ |
|
|
|
|
|
'name': dataArray.name, |
|
|
|
|
|
'mongoID': mongoID, |
|
|
|
|
|
'engine': 'xarray', |
|
|
|
|
|
'dtype': 'dataArray', |
|
|
|
|
|
} |
|
|
|
|
|
} |
|
|
|
|
|
|
|
|
filter = { |
|
|
|
|
|
'year': self.year, |
|
|
|
|
|
'month': self.month, |
|
|
|
|
|
'day': self.day, |
|
|
|
|
|
'shotNum': shotNum, |
|
|
|
|
|
} |
|
|
|
|
|
|
|
|
self.mongoDB['global'].update_one(filter, {"$set": data_label}, upsert=False) |
|
|
|
|
|
|
|
|
for i in stackedDataArray: |
|
|
|
|
|
|
|
|
|
|
|
stackedDataArray_single = i[1].drop('_scanAxis') |
|
|
|
|
|
|
|
|
|
|
|
global_parameters = { |
|
|
|
|
|
'global_parameters.' + key: stackedDataArray_single[key].item() |
|
|
|
|
|
for key in scanAxis |
|
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
|
|
filter.update(global_parameters) |
|
|
|
|
|
|
|
|
|
|
|
mongoID, _ = self.xdb.put(stackedDataArray_single) |
|
|
|
|
|
|
|
|
|
|
|
data_label = { |
|
|
|
|
|
dataArray.name: |
|
|
|
|
|
{ |
|
|
|
|
|
'name': dataArray.name, |
|
|
|
|
|
'mongoID': mongoID, |
|
|
|
|
|
'engine': 'xarray', |
|
|
|
|
|
'dtype': 'dataArray', |
|
|
|
|
|
} |
|
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
|
|
self.mongoDB['global'].update_one(filter, {"$set": data_label}, upsert=False) |
|
|
|
|
|
|
|
|
def _add_data_xarray_dataSet(self, shotNum, dataSet, name): |
|
|
|
|
|
|
|
|
def _add_data_xarray_dataSet(self, shotNum, dataSet, name, scanAxis=None): |
|
|
|
|
|
|
|
|
|
|
|
if scanAxis is None: |
|
|
|
|
|
scanAxis = list(dataSet.coords) |
|
|
|
|
|
|
|
|
|
|
|
dataSet.attrs = self._convert_numpy_type(dataSet.attrs) |
|
|
|
|
|
|
|
|
|
|
|
for key in list(dataSet.data_vars): |
|
|
|
|
|
dataSet[key].attrs = self._convert_numpy_type(dataSet[key].attrs) |
|
|
|
|
|
|
|
|
|
|
|
stackedDataSet = dataSet.stack(_scanAxis=tuple(scanAxis)) |
|
|
|
|
|
stackedDataSet = stackedDataSet.groupby('_scanAxis') |
|
|
|
|
|
|
|
|
filter = { |
|
|
filter = { |
|
|
'year': self.year, |
|
|
'year': self.year, |
|
|
'month': self.month, |
|
|
'month': self.month, |
|
@ -106,32 +192,42 @@ class MongoDB: |
|
|
'shotNum': shotNum, |
|
|
'shotNum': shotNum, |
|
|
} |
|
|
} |
|
|
|
|
|
|
|
|
dataSet.attrs = self._convert_numpy_type(dataSet.attrs) |
|
|
|
|
|
|
|
|
|
|
|
for key in list(dataSet.data_vars): |
|
|
|
|
|
dataSet[key].attrs = self._convert_numpy_type(dataSet[key].attrs) |
|
|
|
|
|
print(key) |
|
|
|
|
|
|
|
|
for i in stackedDataSet: |
|
|
|
|
|
|
|
|
|
|
|
stackedDataSet_single = i[1].drop('_scanAxis') |
|
|
|
|
|
|
|
|
|
|
|
global_parameters = { |
|
|
|
|
|
'global_parameters.' + key: stackedDataSet_single[key].item() |
|
|
|
|
|
for key in scanAxis |
|
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
|
|
filter.update(global_parameters) |
|
|
|
|
|
|
|
|
mongoID, _ = self.xdb.put(dataSet) |
|
|
|
|
|
|
|
|
mongoID, _ = self.xdb.put(dataSet) |
|
|
|
|
|
|
|
|
data_label = { |
|
|
|
|
|
name: |
|
|
|
|
|
{ |
|
|
|
|
|
'name': name, |
|
|
|
|
|
'mongoID': mongoID, |
|
|
|
|
|
'engine': 'xarray', |
|
|
|
|
|
'dtype': 'dataSet', |
|
|
|
|
|
} |
|
|
|
|
|
} |
|
|
|
|
|
|
|
|
data_label = { |
|
|
|
|
|
name: |
|
|
|
|
|
{ |
|
|
|
|
|
'name': name, |
|
|
|
|
|
'mongoID': mongoID, |
|
|
|
|
|
'engine': 'xarray', |
|
|
|
|
|
'dtype': 'dataSet', |
|
|
|
|
|
} |
|
|
|
|
|
} |
|
|
|
|
|
|
|
|
self.mongoDB['global'].update_one(filter, {"$set": data_label}, upsert=False) |
|
|
|
|
|
|
|
|
self.mongoDB['global'].update_one(filter, {"$set": data_label}, upsert=False) |
|
|
|
|
|
|
|
|
def _add_data_additional(self, shotNum, data, name): |
|
|
|
|
|
|
|
|
def _add_data_additional(self, shotNum, runNum, data, name): |
|
|
|
|
|
|
|
|
|
|
|
if runNum is None: |
|
|
|
|
|
runNum = 0 |
|
|
|
|
|
|
|
|
filter = { |
|
|
filter = { |
|
|
'year': self.year, |
|
|
'year': self.year, |
|
|
'month': self.month, |
|
|
'month': self.month, |
|
|
'day': self.day, |
|
|
'day': self.day, |
|
|
'shotNum': shotNum, |
|
|
'shotNum': shotNum, |
|
|
|
|
|
'runNum': runNum, |
|
|
} |
|
|
} |
|
|
|
|
|
|
|
|
mongoID = self.mongoDB.additional.insert_one(data).inserted_id |
|
|
mongoID = self.mongoDB.additional.insert_one(data).inserted_id |
|
@ -148,21 +244,21 @@ class MongoDB: |
|
|
|
|
|
|
|
|
self.mongoDB['global'].update_one(filter, {"$set": data_label}, upsert=False) |
|
|
self.mongoDB['global'].update_one(filter, {"$set": data_label}, upsert=False) |
|
|
|
|
|
|
|
|
def add_data(self, shotNum, data, date=None, name=None, engine='normal'): |
|
|
|
|
|
|
|
|
def add_data(self, shotNum, data, runNum=None, date=None, name=None, engine='normal'): |
|
|
if not date is None: |
|
|
if not date is None: |
|
|
self.set_date(date) |
|
|
self.set_date(date) |
|
|
|
|
|
|
|
|
if engine == 'normal': |
|
|
if engine == 'normal': |
|
|
self._add_data_normal(shotNum=shotNum, data=data) |
|
|
|
|
|
|
|
|
self._add_data_normal(shotNum=shotNum, runNum=runNum, data=data) |
|
|
elif engine == 'xarray': |
|
|
elif engine == 'xarray': |
|
|
if isinstance(data, type(xr.Dataset())): |
|
|
if isinstance(data, type(xr.Dataset())): |
|
|
self._add_data_xarray_dataSet(shotNum=shotNum, data=data, name=name) |
|
|
self._add_data_xarray_dataSet(shotNum=shotNum, data=data, name=name) |
|
|
else: |
|
|
else: |
|
|
self._add_data_xarray_dataArray(shotNum=shotNum, data=data, name=name) |
|
|
self._add_data_xarray_dataArray(shotNum=shotNum, data=data, name=name) |
|
|
elif engine == 'additional': |
|
|
elif engine == 'additional': |
|
|
self._add_data_additional(shotNum=shotNum, data=data, name=name) |
|
|
|
|
|
|
|
|
self._add_data_additional(shotNum=shotNum, runNum=runNum, data=data, name=name) |
|
|
|
|
|
|
|
|
def read_global(self, shotNum, date=None): |
|
|
|
|
|
|
|
|
def read_global_single(self, shotNum, runNum, date=None): |
|
|
|
|
|
|
|
|
if not date is None: |
|
|
if not date is None: |
|
|
self.set_date(date) |
|
|
self.set_date(date) |
|
@ -172,6 +268,7 @@ class MongoDB: |
|
|
'month': self.month, |
|
|
'month': self.month, |
|
|
'day': self.day, |
|
|
'day': self.day, |
|
|
'shotNum': shotNum, |
|
|
'shotNum': shotNum, |
|
|
|
|
|
'runNum': runNum, |
|
|
} |
|
|
} |
|
|
|
|
|
|
|
|
return self.mongoDB['global'].find_one(filter) |
|
|
return self.mongoDB['global'].find_one(filter) |
|
@ -182,18 +279,18 @@ class MongoDB: |
|
|
if engine == 'additional': |
|
|
if engine == 'additional': |
|
|
return self.mongoDB.additional.find_one({'_id': mongoID}) |
|
|
return self.mongoDB.additional.find_one({'_id': mongoID}) |
|
|
|
|
|
|
|
|
def load_data_single(self, shotNum=None, globalDict=None, date=None, field=None): |
|
|
|
|
|
|
|
|
def load_data_single(self, shotNum=None, runNum=None, globalDict=None, date=None, field=None): |
|
|
|
|
|
|
|
|
if not date is None: |
|
|
if not date is None: |
|
|
self.set_date(date) |
|
|
self.set_date(date) |
|
|
|
|
|
|
|
|
if globalDict is None: |
|
|
if globalDict is None: |
|
|
globalDict = self.read_global(shotNum) |
|
|
|
|
|
|
|
|
globalDict = self.read_global_single(shotNum=shotNum, runNum=runNum) |
|
|
|
|
|
|
|
|
if field is None: |
|
|
if field is None: |
|
|
field = globalDict |
|
|
field = globalDict |
|
|
|
|
|
|
|
|
res = {} |
|
|
|
|
|
|
|
|
res = field |
|
|
|
|
|
|
|
|
for key in field: |
|
|
for key in field: |
|
|
if isinstance(globalDict[key], dict) and ('mongoID' in globalDict[key]): |
|
|
if isinstance(globalDict[key], dict) and ('mongoID' in globalDict[key]): |
|
|