liuyu/qlib
1
0
Fork 0
mirror of https://github.com/microsoft/qlib.git synced 2026-07-05 03:50:57 +08:00
Code Activity

Update seperate dataframe

Browse Source
This commit is contained in:
Young
2021-11-01 00:51:21 +08:00
parent 31e9d529de
commit 82f8ff9066
1 changed files with 26 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

36
qlib/contrib/data/utils/sepdf.py
View File

@@ -75,6 +75,10 @@ class SepDataFrame:
def copy(self, *args, **kwargs):
return self.apply_each("copy", True, *args, **kwargs)
def _update_join(self):
if self.join not in self:
self.join = next(iter(self._df_dict.keys()))
def __getitem__(self, item):
return self._df_dict[item]
@@ -82,9 +86,16 @@ class SepDataFrame:
# TODO: consider the join behavior
self._df_dict[item] = df
def __delitem__(self, item: str):
del self._df_dict[item]
self._update_join()
def __contains__(self, item):
return item in self._df_dict
def __len__(self):
return len(self._df_dict[self.join])
def droplevel(self, *args, **kwargs):
raise NotImplementedError(f"Please implement the `droplevel` method")
@@ -125,21 +136,26 @@ class SDFLoc:
return self._sdf[args]
elif isinstance(args, (tuple, list)):
new_df_dict = {k: self._sdf[k] for k in args}
return SepDataFrame(new_df_dict, join=self.join if self.join in args else args[0])
return SepDataFrame(new_df_dict, join=self.join if self.join in args else args[0], skip_align=True)
else:
raise NotImplementedError(f"This type of input is not supported")
elif self.axis == 0:
return SepDataFrame({k: df.loc(axis=0)[args] for k, df in self._sdf._df_dict.items()}, join=self.join)
return SepDataFrame(
{k: df.loc(axis=0)[args] for k, df in self._sdf._df_dict.items()}, join=self.join, skip_align=True
)
else:
ax0, *ax1 = args
if len(ax1) == 0:
ax1 = None
df = self._sdf
if ax1 is not None:
df = df.loc(axis=1)[ax1]
if ax0 is not None:
df = df.loc(axis=0)[ax0]
return df
if isinstance(args, tuple):
ax0, *ax1 = args
if len(ax1) == 0:
ax1 = None
if ax1 is not None:
df = df.loc(axis=1)[ax1]
if ax0 is not None:
df = df.loc(axis=0)[ax0]
return df
else:
return df.loc(axis=0)[args]
# Patch pandas DataFrame