|
@ -16,8 +16,19 @@ npArrayType = type(np.array([0])) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class MongoDB: |
|
|
class MongoDB: |
|
|
|
|
|
"""A class for communicate with our MongoDB. |
|
|
|
|
|
""" |
|
|
|
|
|
|
|
|
def __init__(self, mongoClient, mongoDB, date=None) -> None: |
|
|
def __init__(self, mongoClient, mongoDB, date=None) -> None: |
|
|
|
|
|
"""Initialize the class with given handle to our MongoDB client and database. |
|
|
|
|
|
|
|
|
|
|
|
:param mongoClient: The handle to MongoDB client |
|
|
|
|
|
:type mongoClient: pymongo MongoDB client or other MongoDB client |
|
|
|
|
|
:param mongoDB: The handle to MongoDB database |
|
|
|
|
|
:type mongoDB: pymongo MongoDB database or other MongoDB database |
|
|
|
|
|
:param date: The date of data ('YYYY/MM/DD'), can be setted later, defaults to None |
|
|
|
|
|
:type date: str, optional |
|
|
|
|
|
""" |
|
|
self.mongoClient = mongoClient |
|
|
self.mongoClient = mongoClient |
|
|
self.mongoDB = mongoDB |
|
|
self.mongoDB = mongoDB |
|
|
self.xdb = xarray_mongodb.XarrayMongoDB(mongoDB) |
|
|
self.xdb = xarray_mongodb.XarrayMongoDB(mongoDB) |
|
@ -27,6 +38,13 @@ class MongoDB: |
|
|
self.set_date(date) |
|
|
self.set_date(date) |
|
|
|
|
|
|
|
|
def _convert_numpy_type(self, data): |
|
|
def _convert_numpy_type(self, data): |
|
|
|
|
|
"""Convert from numpy type to normal python type. |
|
|
|
|
|
|
|
|
|
|
|
:param data: The data need to be converted |
|
|
|
|
|
:type data: numpy data type |
|
|
|
|
|
:return: The converted data |
|
|
|
|
|
:rtype: normal python data type |
|
|
|
|
|
""" |
|
|
for key in data: |
|
|
for key in data: |
|
|
typeKey = type(data[key]) |
|
|
typeKey = type(data[key]) |
|
|
if typeKey in npTypeDict: |
|
|
if typeKey in npTypeDict: |
|
@ -41,6 +59,15 @@ class MongoDB: |
|
|
return data |
|
|
return data |
|
|
|
|
|
|
|
|
def _break_dataSet(self, dataSet, scanAxis=None): |
|
|
def _break_dataSet(self, dataSet, scanAxis=None): |
|
|
|
|
|
"""Stack the scan axes of data |
|
|
|
|
|
|
|
|
|
|
|
:param dataSet: The xarray DataSet or DataArray stored the data |
|
|
|
|
|
:type dataSet: xarray DataSet or DataArray |
|
|
|
|
|
:param scanAxis: a list of the name of scan axes, defaults to None |
|
|
|
|
|
:type scanAxis: list or array like, optional |
|
|
|
|
|
:return: The stacked xarray DataSet or DataArray stored the data |
|
|
|
|
|
:rtype: xarray DataSet or DataArray |
|
|
|
|
|
""" |
|
|
|
|
|
|
|
|
if scanAxis is None: |
|
|
if scanAxis is None: |
|
|
scanAxis = dataSet.attrs['scanAxis'] |
|
|
scanAxis = dataSet.attrs['scanAxis'] |
|
@ -52,12 +79,26 @@ class MongoDB: |
|
|
return stackedDataArray |
|
|
return stackedDataArray |
|
|
|
|
|
|
|
|
def set_date(self, date): |
|
|
def set_date(self, date): |
|
|
|
|
|
"""Set the date of data |
|
|
|
|
|
|
|
|
|
|
|
:param date: The date of data ('YYYY/MM/DD') |
|
|
|
|
|
:type date: str |
|
|
|
|
|
""" |
|
|
date = date.split("/") |
|
|
date = date.split("/") |
|
|
self.year = int(date[0]) |
|
|
self.year = int(date[0]) |
|
|
self.month = int(date[1]) |
|
|
self.month = int(date[1]) |
|
|
self.day = int(date[2]) |
|
|
self.day = int(date[2]) |
|
|
|
|
|
|
|
|
def create_global(self, shotNum, dataSet=None, date=None): |
|
|
def create_global(self, shotNum, dataSet=None, date=None): |
|
|
|
|
|
"""Creat a the global document in MongoDB |
|
|
|
|
|
|
|
|
|
|
|
:param shotNum: The shot number |
|
|
|
|
|
:type shotNum: str |
|
|
|
|
|
:param dataSet: The xarray DataSet stored the global parameters, defaults to None |
|
|
|
|
|
:type dataSet: xarray DataSet, optional |
|
|
|
|
|
:param date: the date of the data, defaults to None |
|
|
|
|
|
:type date: str, optional |
|
|
|
|
|
""" |
|
|
if not date is None: |
|
|
if not date is None: |
|
|
self.set_date(date) |
|
|
self.set_date(date) |
|
|
|
|
|
|
|
@ -133,6 +174,15 @@ class MongoDB: |
|
|
self.mongoDB['global'].insert_one(data) |
|
|
self.mongoDB['global'].insert_one(data) |
|
|
|
|
|
|
|
|
def _add_data_normal(self, shotNum, runNum, data): |
|
|
def _add_data_normal(self, shotNum, runNum, data): |
|
|
|
|
|
"""Write the data directly to the global document |
|
|
|
|
|
|
|
|
|
|
|
:param shotNum: The shot number |
|
|
|
|
|
:type shotNum: str |
|
|
|
|
|
:param runNum: The run number |
|
|
|
|
|
:type runNum: int |
|
|
|
|
|
:param data: The data to be written |
|
|
|
|
|
:type data: normal python data type |
|
|
|
|
|
""" |
|
|
|
|
|
|
|
|
if runNum is None: |
|
|
if runNum is None: |
|
|
runNum = 0 |
|
|
runNum = 0 |
|
@ -148,6 +198,17 @@ class MongoDB: |
|
|
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, name=None, scanAxis=None): |
|
|
def _add_data_xarray_dataArray(self, shotNum, dataArray, name=None, scanAxis=None): |
|
|
|
|
|
"""Write the data in a type of xarray DataArray to the MongoDb. |
|
|
|
|
|
|
|
|
|
|
|
:param shotNum: The shot number |
|
|
|
|
|
:type shotNum: str |
|
|
|
|
|
:param dataArray: The xarray DataArray to be written |
|
|
|
|
|
:type dataArray: xarray DataArray |
|
|
|
|
|
:param name: The name of the DataArray, defaults to None |
|
|
|
|
|
:type name: str, optional |
|
|
|
|
|
:param scanAxis: The scan axes of the data, defaults to None |
|
|
|
|
|
:type scanAxis: array like, optional |
|
|
|
|
|
""" |
|
|
|
|
|
|
|
|
if scanAxis is None: |
|
|
if scanAxis is None: |
|
|
scanAxis = list(dataArray.coords) |
|
|
scanAxis = list(dataArray.coords) |
|
@ -190,6 +251,17 @@ 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_xarray_dataSet(self, shotNum, dataSet, name, scanAxis=None): |
|
|
def _add_data_xarray_dataSet(self, shotNum, dataSet, name, scanAxis=None): |
|
|
|
|
|
"""Write the data in a type of xarray DataSet to the MongoDb. |
|
|
|
|
|
|
|
|
|
|
|
:param shotNum: The shot number |
|
|
|
|
|
:type shotNum: str |
|
|
|
|
|
:param dataSet: The xarray DataSet to be written |
|
|
|
|
|
:type dataSet: xarray DataSet |
|
|
|
|
|
:param name: The name of the DataArray, defaults to None |
|
|
|
|
|
:type name: str, optional |
|
|
|
|
|
:param scanAxis: The scan axes of the data, defaults to None |
|
|
|
|
|
:type scanAxis: array like, optional |
|
|
|
|
|
""" |
|
|
|
|
|
|
|
|
if scanAxis is None: |
|
|
if scanAxis is None: |
|
|
scanAxis = list(dataSet.coords) |
|
|
scanAxis = list(dataSet.coords) |
|
@ -235,6 +307,17 @@ 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_additional(self, shotNum, runNum, data, name): |
|
|
def _add_data_additional(self, shotNum, runNum, data, name): |
|
|
|
|
|
"""Write the data in an additional document |
|
|
|
|
|
|
|
|
|
|
|
:param shotNum: The shot number |
|
|
|
|
|
:type shotNum: str |
|
|
|
|
|
:param runNum: The run number |
|
|
|
|
|
:type runNum: int |
|
|
|
|
|
:param data: The data to be written |
|
|
|
|
|
:type data: normal python data type |
|
|
|
|
|
:param name: The name of the data |
|
|
|
|
|
:type name: str |
|
|
|
|
|
""" |
|
|
|
|
|
|
|
|
if runNum is None: |
|
|
if runNum is None: |
|
|
runNum = 0 |
|
|
runNum = 0 |
|
@ -262,6 +345,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, runNum=None, date=None, name=None, engine='normal'): |
|
|
def add_data(self, shotNum, data, runNum=None, date=None, name=None, engine='normal'): |
|
|
|
|
|
"""Write a new data to MongoDB |
|
|
|
|
|
|
|
|
|
|
|
:param shotNum: The shot number |
|
|
|
|
|
:type shotNum: str |
|
|
|
|
|
:param data: The data needs to be written |
|
|
|
|
|
:type data: standard python data type, numpy type, xarray DataArray or xarray DataSet |
|
|
|
|
|
:param runNum: The run number, defaults to None |
|
|
|
|
|
:type runNum: int, optional |
|
|
|
|
|
:param date: The date of the data ('YYYY/MM/DD'), defaults to None |
|
|
|
|
|
:type date: str, optional |
|
|
|
|
|
:param name: The name of the data, defaults to None |
|
|
|
|
|
:type name: str, optional |
|
|
|
|
|
:param engine: The engine for different types of the data, defaults to 'normal' |
|
|
|
|
|
:type engine: str, optional |
|
|
|
|
|
""" |
|
|
if not date is None: |
|
|
if not date is None: |
|
|
self.set_date(date) |
|
|
self.set_date(date) |
|
|
|
|
|
|
|
@ -276,6 +374,17 @@ class MongoDB: |
|
|
self._add_data_additional(shotNum=shotNum, runNum=runNum, data=data, name=name) |
|
|
self._add_data_additional(shotNum=shotNum, runNum=runNum, data=data, name=name) |
|
|
|
|
|
|
|
|
def read_global_single(self, shotNum, runNum, date=None): |
|
|
def read_global_single(self, shotNum, runNum, date=None): |
|
|
|
|
|
"""Read the global document of specified shot and run from MongoDB. |
|
|
|
|
|
|
|
|
|
|
|
:param shotNum: The shot number |
|
|
|
|
|
:type shotNum: str |
|
|
|
|
|
:param runNum: The run number |
|
|
|
|
|
:type runNum: int |
|
|
|
|
|
:param date: The date of the data ('YYYY/MM/DD'), defaults to None |
|
|
|
|
|
:type date: str, optional |
|
|
|
|
|
:return: The global document |
|
|
|
|
|
:rtype: dict |
|
|
|
|
|
""" |
|
|
|
|
|
|
|
|
if not date is None: |
|
|
if not date is None: |
|
|
self.set_date(date) |
|
|
self.set_date(date) |
|
@ -291,6 +400,15 @@ class MongoDB: |
|
|
return self.mongoDB['global'].find_one(filter) |
|
|
return self.mongoDB['global'].find_one(filter) |
|
|
|
|
|
|
|
|
def read_global_all(self, shotNum, date=None): |
|
|
def read_global_all(self, shotNum, date=None): |
|
|
|
|
|
"""Read the global document of all runs in the specified shot from MongoDB, and extract the scan axes. |
|
|
|
|
|
|
|
|
|
|
|
:param shotNum: The shot number |
|
|
|
|
|
:type shotNum: str |
|
|
|
|
|
:param date: The date of the data ('YYYY/MM/DD'), defaults to None |
|
|
|
|
|
:type date: str, optional |
|
|
|
|
|
:return: The global document |
|
|
|
|
|
:rtype: dict |
|
|
|
|
|
""" |
|
|
|
|
|
|
|
|
from xarray.core.utils import equivalent |
|
|
from xarray.core.utils import equivalent |
|
|
|
|
|
|
|
@ -379,12 +497,37 @@ class MongoDB: |
|
|
return result |
|
|
return result |
|
|
|
|
|
|
|
|
def _load_data_single(self, mongoID, engine): |
|
|
def _load_data_single(self, mongoID, engine): |
|
|
|
|
|
"""load the document according to given _ID |
|
|
|
|
|
|
|
|
|
|
|
:param mongoID: the ID of document in MongoDB |
|
|
|
|
|
:type mongoID: MongoDB ID object |
|
|
|
|
|
:param engine: _description_ |
|
|
|
|
|
:type engine: _type_ |
|
|
|
|
|
:return: The engine for different types of the data |
|
|
|
|
|
:rtype: str |
|
|
|
|
|
""" |
|
|
|
|
|
|
|
|
if engine == 'xarray': |
|
|
if engine == 'xarray': |
|
|
return self.xdb.get(mongoID) |
|
|
return self.xdb.get(mongoID) |
|
|
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, runNum=None, globalDict=None, date=None, field=None): |
|
|
def load_data_single(self, shotNum=None, runNum=None, globalDict=None, date=None, field=None): |
|
|
|
|
|
"""go through the given global document and find all the MongoDB object, then replace them with the document they linked to. |
|
|
|
|
|
|
|
|
|
|
|
:param shotNum: The shot number, defaults to None |
|
|
|
|
|
:type shotNum: str, optional |
|
|
|
|
|
:param runNum: The run number, defaults to None |
|
|
|
|
|
:type runNum: int, optional |
|
|
|
|
|
:param globalDict: The given global document, defaults to None |
|
|
|
|
|
:type globalDict: dict, optional |
|
|
|
|
|
:param date: The date of the data ('YYYY/MM/DD'), defaults to None |
|
|
|
|
|
:type date: str, optional |
|
|
|
|
|
:param field: If not None, the function will only go through these keys, defaults to None |
|
|
|
|
|
:type field: array like, optional |
|
|
|
|
|
:return: The document with loaded data |
|
|
|
|
|
:rtype: dict |
|
|
|
|
|
""" |
|
|
|
|
|
|
|
|
if not date is None: |
|
|
if not date is None: |
|
|
self.set_date(date) |
|
|
self.set_date(date) |
|
@ -411,4 +554,5 @@ class MongoDB: |
|
|
return res |
|
|
return res |
|
|
|
|
|
|
|
|
def load_data(self, shotNum=None, globalDict=None, date=None, field=None): |
|
|
def load_data(self, shotNum=None, globalDict=None, date=None, field=None): |
|
|
|
|
|
# load all the data of specified shot |
|
|
pass |
|
|
pass |