analyseScript/DataContainer/MongoDB.py

212 lines
6.3 KiB
Python
Raw Normal View History

2023-06-09 18:59:56 +02:00
import pymongo
import xarray_mongodb
import bson
import builtins
2023-06-12 12:52:07 +02:00
import xarray as xr
2023-06-09 18:59:56 +02:00
from ToolFunction.ToolFunction import get_date
npTypeDict = {v: getattr(builtins, k) for k, v in np.sctypeDict.items() if k in vars(builtins)}
npArrayType = type(np.array([0]))
class MongoDB:
def __init__(self, mongoClient, mongoDB, date=None) -> None:
self.mongoClient = mongoClient
self.mongoDB = mongoDB
self.xdb = xarray_mongodb.XarrayMongoDB(mongoDB)
if date is None:
date= get_date()
self.set_date(date)
def _convert_numpy_type(self, data):
for key in data:
typeKey = type(data[key])
if typeKey in npTypeDict:
data[key] = data[key].item()
elif typeKey == npArrayType:
data[key] = data[key].tolist()
else:
try:
data[key] = data[key].item()
except:
pass
return data
def set_date(self, date):
date = date.split("/")
self.year = int(date[0])
self.month = int(date[1])
self.day = int(date[2])
def create_global(self, shotNum, dataSet=None, date=None):
if not date is None:
self.set_date(date)
data = {
'year': self.year,
'month': self.month,
'day': self.day,
'shotNum': shotNum,
'global_parameters' : {},
}
global_parameters = self._convert_numpy_type(dataSet.attrs)
if not dataSet is None:
data['global_parameters'].update(global_parameters)
data = self._convert_numpy_type(data)
self.mongoDB['global'].insert_one(data)
def _add_data_normal(self, shotNum, data):
filter = {
'year': self.year,
'month': self.month,
'day': self.day,
'shotNum': shotNum,
}
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,
}
dataArray.attrs = self._convert_numpy_type(dataArray.attrs)
mongoID, _ = self.xdb.put(dataArray)
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):
filter = {
'year': self.year,
'month': self.month,
'day': self.day,
'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)
mongoID, _ = self.xdb.put(dataSet)
data_label = {
name:
{
'name': name,
'mongoID': mongoID,
'engine': 'xarray',
'dtype': 'dataSet',
}
}
self.mongoDB['global'].update_one(filter, {"$set": data_label}, upsert=False)
2023-06-12 12:52:07 +02:00
def _add_data_additional(self, shotNum, data, name):
filter = {
'year': self.year,
'month': self.month,
'day': self.day,
'shotNum': shotNum,
}
mongoID = self.mongoDB.additional.insert_one(data).inserted_id
data_label = {
name:
{
'name': name,
'mongoID': mongoID,
'engine': 'additional',
'dtype': 'dict',
}
}
self.mongoDB['global'].update_one(filter, {"$set": data_label}, upsert=False)
2023-06-09 18:59:56 +02:00
def add_data(self, shotNum, data, date=None, name=None, engine='normal'):
2023-06-09 19:34:37 +02:00
if not date is None:
self.set_date(date)
if engine == 'normal':
self._add_data_normal(shotNum=shotNum, data=data)
elif engine == 'xarray':
2023-06-12 12:52:07 +02:00
if isinstance(data, type(xr.Dataset())):
2023-06-09 19:34:37 +02:00
self._add_data_xarray_dataSet(shotNum=shotNum, data=data, name=name)
else:
self._add_data_xarray_dataArray(shotNum=shotNum, data=data, name=name)
2023-06-12 12:52:07 +02:00
elif engine == 'additional':
self._add_data_additional(shotNum=shotNum, data=data, name=name)
def read_global(self, shotNum, date=None):
if not date is None:
self.set_date(date)
filter = {
'year': self.year,
'month': self.month,
'day': self.day,
'shotNum': shotNum,
}
return self.mongoDB['global'].find_one(filter)
def _load_data_single(self, mongoID, engine):
if engine == 'xarray':
return self.xdb.get(mongoID)
if engine == 'additional':
return self.mongoDB.additional.find_one({'_id': mongoID})
2023-06-09 18:59:56 +02:00
2023-06-12 12:52:07 +02:00
def load_data_single(self, shotNum=None, globalDict=None, date=None, field=None):
if not date is None:
self.set_date(date)
if globalDict is None:
globalDict = self.read_global(shotNum)
if field is None:
field = globalDict
res = {}
for key in field:
if isinstance(globalDict[key], dict) and ('mongoID' in globalDict[key]):
mongoID = globalDict[key]['mongoID']
engine = globalDict[key]['engine']
res.update(
{
key: self._load_data_single(mongoID=mongoID, engine=engine)
}
)
return res
def load_data(self, shotNum=None, globalDict=None, date=None, field=None):
2023-06-09 18:59:56 +02:00
pass