1
0
mirror of https://github.com/microsoft/qlib.git synced 2026-07-07 04:50:56 +08:00
Files
qlib/qlib/finco/context.py
you-n-g 1c9841b15e Connect TrainTask & Rolling & DDG-DA (#1599)
* Connect train task to ddg-da & rolling

* Pylint & black formatting

* Formatting
2023-07-17 09:58:58 +08:00

98 lines
3.2 KiB
Python

from dataclasses import dataclass, field
import copy
from pathlib import Path
from typing import Optional, List
from qlib.finco.log import FinCoLog
from qlib.typehint import Literal
from qlib.finco.utils import similarity
@dataclass
class Design:
plan: str
classes: str
decision: str
@dataclass
class Exp:
"""Experiment"""
# compoments
dataset: Optional[Design] = None
datahandler: Optional[Design] = None
model: Optional[Design] = None
record: Optional[Design] = None
strategy: Optional[Design] = None
backtest: Optional[Design] = None
# basic
template: Optional[Path] = None
# rolling strategy. None indicates no rolling
rolling: Optional[Literal["base", "ddgda"]] = None
@dataclass
class StructContext:
"""Part of the context have clear meaning and structure, so they will be saved here and can be easily retrieved and understood"""
# TODO: move more content in WorkflowContextManager.context to here
workspace: Path
exp_list: List[Exp] = field(default_factory=list) # the planned experiments
class WorkflowContextManager:
"""Context Manager stores the context of the workflow"""
"""All context are key value pairs which saves the input, output and status of the whole workflow"""
def __init__(self, workspace: Path) -> None:
self.context = {}
self.logger = FinCoLog()
# this context is public
self.struct_context = StructContext(workspace) # TODO: move more content in context to here
self.set_context("workspace", workspace) # TODO: remove me
def set_context(self, key, value):
if key in self.context:
self.logger.warning("The key already exists in the context, the value will be overwritten")
self.context[key] = value
def get_context(self, key):
# NOTE: if the key doesn't exist, return None. In the future, we may raise an error to detect abnormal behavior
if key not in self.context:
self.logger.warning("The key doesn't exist in the context")
return None
return self.context[key]
def update_context(self, key, new_value):
# NOTE: if the key doesn't exist, return None. In the future, we may raise an error to detect abnormal behavior
if key not in self.context:
self.logger.warning("The key doesn't exist in the context")
self.context.update({key: new_value})
def get_all_context(self):
"""return a deep copy of the context"""
"""TODO: do we need to return a deep copy?"""
return copy.deepcopy(self.context)
def retrieve(self, query: str) -> dict:
if query in self.context.keys():
return {query: self.context.get(query)}
# Note: retrieve information from context by string similarity maybe abandon in future
scores = {}
for k, v in self.context.items():
scores.update({k: max(similarity(query, k), similarity(query, v))})
max_score_key = max(scores, key=scores.get)
return {max_score_key: self.context.get(max_score_key)}
def clear(self, reserve: list = None):
if reserve is None:
reserve = []
_context = {k: self.get_context(k) for k in reserve}
self.context = _context