jacksung-dev 0.0.4.15__py3-none-any.whl

jacksung/ai/latex_tool.py

@@ -0,0 +1,199 @@
+import os
+from openai import OpenAI
+from tqdm import tqdm
+from jacksung.utils.time import Stopwatch, get_time_str
+
+
+def get_en_polish_prompt(text, prompt_type='polish'):
+    polish_prompt = \
+        fr'''
+        # Rewrite the text in an academic writing style, using more appropriate vocabulary and sentence structure while keeping the original meaning unchanged:
+        - Make sure the rewritten version conveys the same information and intention as the original text.
+        - Please output the rewritten text directly in latex format, without including the original text, thinking logic, comments, explanations, etc.
+        - Do not output any control commands that do not exist in the input content (such as \documentclass, \begin, \end, etc.), just use latex format to output mathematical formulas, symbols, references, or other latex instructions contained in the input content.
+        - Note that special symbols and formulas are output in latex format, not directly output special characters.
+        - Only when the input content contains control codes such as \par, the code needs to be added to the corresponding position of the output content.
+        - If the input content only contains code and does not contain any substantial text content, the input content is directly output without any changes. Be careful not to miss symbols such as brackets.
+        - Make sure that the output content can be compiled normally after replacing the input content in the original document.
+        The following is the input content:
+        {text}
+        '''
+    check_prompt = \
+        fr'''
+        # Correct grammatical errors and misspelled words in the input text while preserving the original meaning and format:
+        - If the text is free of grammatical errors and spelling mistakes, output the original text without making any modifications
+        - Ensure the corrected version maintains the exact same information and intent as the original text
+        - Output only the corrected text directly without any additional content
+        - Preserve all mathematical formulas, symbols, and special formatting exactly as input
+        - Maintain the original document structure and formatting commands
+        - Only modify actual grammatical errors and misspelled words
+        - If the input contains LaTeX code, preserve it exactly and only correct text outside code blocks
+        - Do not alter technical terms, proper nouns, or specialized vocabulary
+        - If no errors are detected, output the original text unchanged
+        - Do not include any explanations, comments, or thinking process in the output
+
+        Input text:
+        {text}
+        '''
+    if prompt_type == 'polish':
+        return polish_prompt
+    elif prompt_type == 'check':
+        return check_prompt
+    else:
+        raise Exception(rf'Unknown prompt type {prompt_type}, please specify "polish" or "check"')
+
+
+def get_cn_polish_prompt(text, prompt_type='polish'):
+    polish_prompt = \
+        fr'''
+        # 用学术写作风格重写下面的文本,在保持原本涵义不变的情况下使用更合适的词汇和句子结构:
+        - 确保改写后的版本传达的信息和意图与原文相同
+        - 请直接以latex格式输出重写后的文本，不需要包含原文、思考逻辑、注释、解释说明等其他内容。
+        - 不需要输出任何输入内容中不存在的控制命令（如\documentclass、\begin、\end等），只需要使用latex格式输出数学公式、符号、引用或者输入内容中所包含的其他latex指令。
+        - 注意特殊符号和公式以latex格式输出，而不是直接输出特殊字符。
+        - 仅在输入内容中包含\par类似的控制性代码时，在输出内容对应位置需要添加该代码。
+        - 如果输入内容仅包含代码，不包含任何实质性的文本内容，则直接将输入内容不做任何改动输出。注意不要遗漏括号等符号。
+        - 确保输出内容在原文档中替换输入内容后能够正常编译通过。
+        - 重要：你需要确定输入内容所使用的语言，然后使用相同语言进行输出，以保证输入和输出语言一致。
+        以下为输入内容:
+        {text}
+        '''
+    check_prompt = \
+        fr'''
+        # 修正输入文本中的语法错误和单词拼写错误，同时保持原意和格式不变：\
+        - 如果如数的文本没有语法错误和单词拼写则原文输出，不做任何改动
+        - 确保修正后的文本与原文的信息和意图完全相同
+        - 请直接输出修正后的文本，不需要包含原文、思考逻辑、注释、解释说明等其他内容
+        - 不需要输出任何输入内容中不存在的控制命令（如\documentclass、\begin、\end等），只需要使用latex格式输出数学公式、符号、引用或者输入内容中所包含的其他latex指令
+        - 注意特殊符号和公式以latex格式输出，而不是直接输出特殊字符
+        - 仅在输入内容中包含\par类似的控制性代码时，在输出内容对应位置需要添加该代码
+        - 如果输入内容仅包含代码，不包含任何实质性的文本内容，则直接将输入内容不做任何改动输出。注意不要遗漏括号等符号
+        - 确保输出内容在原文档中替换输入内容后能够正常编译通过
+        - 重要：你需要确定输入内容所使用的语言，然后使用相同语言进行输出，以保证输入和输出语言一致
+        以下为输入内容:
+        {text}
+        '''
+    if prompt_type == 'polish':
+        return polish_prompt
+    elif prompt_type == 'check':
+        return check_prompt
+    else:
+        raise Exception(rf'Unknown prompt type {prompt_type}, please specify "polish" or "check"')
+
+
+def merge_content(tex_dir, main_tex):
+    result_tex = ''
+    with open(os.path.join(tex_dir, main_tex), 'r', encoding='utf-8') as f:
+        while True:
+            line = f.readline()
+            if not line:
+                break
+            if line.startswith(r'\input{') or line.startswith(r'\include{'):
+                sub_tex_path = line.split('{')[1].split('}')[0]
+                if not sub_tex_path.endswith('.tex'):
+                    sub_tex_path += r'.tex'
+                result_tex += merge_content(tex_dir, sub_tex_path) + '\n'
+            else:
+                result_tex += line
+    return result_tex
+
+
+class AI:
+    def __init__(self, token, base_url, model_name='deepseek-r1:70b', prompt_type='polish'):
+        self.client = OpenAI(api_key=token, base_url=base_url)
+        self.model_name = model_name
+        self.prompt_type = prompt_type
+
+    def call_ai_polish(self, text, cn_prompt=False, prompt=None):
+        response = self.client.chat.completions.create(
+            model=self.model_name,
+            messages=[
+                {"role": "user",
+                 "content": ((get_cn_polish_prompt(text, self.prompt_type) if cn_prompt else get_en_polish_prompt(
+                     text, self.prompt_type)) if prompt is None else prompt.replace('{content}', text))}
+            ],
+            temperature=0.6,
+            # max_tokens=1024,
+            stream=False
+        )
+        # 逐步接收并处理响应
+        # for chunk in response:
+        #     print(chunk.choices[0].delta.content, end='')
+        # print(response.choices[0].message.content)
+        content = response.choices[0].message.content
+        content = content.split('</think>')[1].strip().replace('\n\n', ' ')
+        if text.startswith(r'\par ') and not content.startswith(r'\par '):
+            print(rf'missing \par in polished text, append it to the beginning of the text.')
+            content = r'\par ' + content
+        if text.startswith(r'\par{') and not content.startswith(r'\par{'):
+            print(rf'missing \par in polished text, append it to the beginning of the text.')
+            content = r'\par{' + content
+        if text.endswith(r'}') and not content.endswith(r'}'):
+            print(r'missing } in polished text, append it to the beginning of the text.')
+            content = content + r'}'
+        return content
+
+
+def polish(main_dir_path, tex_file, server_url, token='Your token here', model_name='deepseek-r1:70b', cn_prompt=False,
+           prompt=None, rewrite_list=(r'\caption{', r'\par ', r'\par{'), skip_part_list=('figure', 'table', 'equation'),
+           ignore_length=100, prompt_type='polish'):
+    st = Stopwatch()
+    ai = AI(token=token, base_url=server_url, model_name=model_name, prompt_type=prompt_type)
+    result_tex = merge_content(main_dir_path, tex_file)
+    new_tex = ''
+    up_flag = False
+    result_split = result_tex.split('\n')
+    for idx, line in enumerate(result_split):
+        spend_count = Stopwatch()
+        line = line.strip()
+        line_up_flag = True
+        if line.startswith('%') or line.startswith('\\') or len(line) < ignore_length:
+            for flag in rewrite_list:
+                if line.startswith(flag):
+                    line_up_flag = False
+                    break
+        else:
+            line_up_flag = False
+
+        for flag in skip_part_list:
+            if line.count(r'\begin{' + flag) > 0:
+                up_flag = True
+                break
+            if line.count(r'\end{' + flag) > 0:
+                up_flag = False
+                break
+
+        if up_flag or line_up_flag:
+            new_tex += line + '\n'
+        else:
+            try:
+                print(rf'当前处理{idx}/{len(result_split)}行,总共用时{st.pinch()},当前时间:{get_time_str()}')
+                print(f'Input[{line[:100]}{"..." if len(line) > 100 else line}]')
+                polish_text = ai.call_ai_polish(line, cn_prompt, prompt)
+                print(f'Polished:[{polish_text[:100]}{"..." if len(polish_text) > 100 else polish_text}]')
+                print(rf'处理结束，耗时{spend_count.pinch()},共改写{len(line)}个字符')
+                new_tex += polish_text + '\n'
+            except Exception as e:
+                tqdm.write(f'**e**{e}')
+                new_tex += line + '\n'
+
+    with open(rf'{main_dir_path}\old.tex', 'w', encoding='utf-8') as f:
+        f.write(result_tex)
+    with open(rf'{main_dir_path}\new.tex', 'w', encoding='utf-8') as f:
+        f.write(new_tex)
+    write_diff(main_dir_path)
+
+
+def write_diff(dir_path):
+    diff_tex = '''\RequirePackage{shellesc}
+    \ShellEscape{pdfLatex new.tex} %编译新文档
+    \ShellEscape{pdfLatex old.tex} %编译新文档
+    \ShellEscape{latexdiff old.tex new.tex > diff_result.tex}
+    \input{diff_result}
+    \documentclass{dummy}'''
+    with open(rf'{dir_path}\diff.tex', 'w', encoding='utf-8') as f:
+        f.write(diff_tex)
+
+
+if __name__ == "__main__":
+    pass

jacksung/ai/metrics.py

@@ -0,0 +1,181 @@
+import os
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.autograd import Variable
+import math
+from pytorch_msssim import ssim
+from torchmetrics.image import PeakSignalNoiseRatio, StructuralSimilarityIndexMeasure
+from torchmetrics import R2Score, PearsonCorrCoef, AUROC
+from torchmetrics.regression import MeanSquaredError
+import importlib
+from einops import rearrange
+import cv2
+
+
+def compute_rmse(da_fc, da_true):
+    error = da_fc - da_true
+    error = error ** 2
+    number = torch.sqrt(error.mean((-2, -1)))
+    return number.mean()
+
+
+class Metrics:
+    def __init__(self):
+        self.psnr = PeakSignalNoiseRatio()
+        self.ssim = StructuralSimilarityIndexMeasure()
+        self.rr = R2Score()
+        self.p = PearsonCorrCoef()
+        self.AUROC = AUROC("binary")
+
+    def mask_nan(self, pred, target):
+        pred = pred.flatten()
+        target = target.flatten()
+        # 生成有效掩码：真实值和预测值均非 NaN 的位置为 True
+        valid_mask = ~(torch.isnan(pred) | torch.isnan(target))
+
+        # 过滤无效元素
+        pred = pred[valid_mask]
+        target = target[valid_mask]
+
+        return pred, target
+
+    def calc_AUROC(self, preds, targets):
+        AUROC = 0
+        for i in range(len(preds)):
+            pred, target = self.mask_nan(preds[i], targets[i])
+            AUROC += self.AUROC(pred, target)
+        self.AUROC.reset()
+        return AUROC / len(preds)
+
+    def calc_psnr(self, preds, targets):
+        psnr = 0
+        for i in range(len(preds)):
+            psnr += self.psnr(rearrange(preds[i], '(b c) h w->b c h w', b=1),
+                              rearrange(targets[i], '(b c) h w->b c h w', b=1))
+        self.psnr.reset()
+        return psnr / len(preds)
+
+    def calc_ssim(self, preds, targets):
+        ssim = 0
+        for i in range(len(preds)):
+            ssim += self.ssim(rearrange(preds[i], '(b c) h w->b c h w', b=1),
+                              rearrange(targets[i], '(b c) h w->b c h w', b=1))
+        self.ssim.reset()
+        return ssim / len(preds)
+
+    def calc_rmse(self, preds, targets):
+        rmse = 0
+        for i in range(len(preds)):
+            rmse += compute_rmse(preds[i], targets[i])
+        return rmse / len(preds)
+
+    def calc_rr(self, preds, targets):
+        rr = 0
+        for i in range(len(preds)):
+            pred, tar = self.mask_nan(preds[i], targets[i])
+            rr += self.rr(pred, tar)
+        self.rr.reset()
+        return rr / len(preds)
+
+    def calc_p(self, preds, targets, exclude_zero=False):
+        p = 0
+        count = 0
+        for i in range(len(preds)):
+            pred, target = self.mask_nan(preds[i], targets[i])
+            if exclude_zero:
+                mask = target != 0
+                pred = pred[mask]
+                target = target[mask]
+            if pred.var() == 0 or target.var() == 0:
+                continue
+            count += 1
+            p += self.p(pred, target)
+        self.p.reset()
+        return p / count
+
+    def print_metrics(self, preds, targets, print_log=True):
+        rr = float(self.calc_rr(preds, targets))
+        p = float(self.calc_p(preds, targets))
+        rmse = float(self.calc_rmse(preds, targets))
+        ssim = float(self.calc_ssim(preds, targets))
+        psnr = float(self.calc_psnr(preds, targets))
+        if print_log:
+            print(rf'p: {p} rr: {rr} rmse: {rmse} ssim: {ssim} psnr: {psnr}')
+        return {'p': p, 'rr': rr, 'rmse': rmse, 'ssim': ssim, 'psnr': psnr}
+
+    def calculate_rain_metrics(self, preds, targets, threshold=0.1):
+        """
+        使用flatten()批量计算多张降雨图的POD、FAR、ACC、CSI指标
+
+        参数:
+            preds: 预测的降雨tensor，形状为[样本数, 高度, 宽度]
+            target: 观测的降雨tensor，形状与preds相同
+            threshold: 降雨事件的阈值，默认0.1mm
+
+        返回:
+            metrics: 包含每个样本指标的字典
+        """
+        # 1. 将每个样本的空间维度展平（[样本数, 高度, 宽度] → [样本数, 像素总数]）
+        preds_flat = preds.flatten(start_dim=1)  # 从第1维开始展平（保留样本维度）
+        targets_flat = targets.flatten(start_dim=1)
+
+        # 2. 二值化（1=有雨，0=无雨）
+        preds_binary = (preds_flat >= threshold).float()
+        targets_binary = (targets_flat >= threshold).float()
+
+        # 3. 计算混淆矩阵元素（按样本维度求和）
+        TP = torch.sum(preds_binary * targets_binary, dim=1)  # 每个样本的TP总和
+        FP = torch.sum(preds_binary * (1 - targets_binary), dim=1)
+        TN = torch.sum((1 - preds_binary) * (1 - targets_binary), dim=1)
+        FN = torch.sum((1 - preds_binary) * targets_binary, dim=1)
+
+        # 4. 计算指标（处理分母为0的情况）
+        POD = TP / (TP + FN)
+        FAR = FP / (TP + FP)
+        ACC = (TP + TN) / (TP + FP + TN + FN)
+        CSI = TP / (TP + FP + FN)
+
+        # 5. 标记无效值为NaN
+        POD = torch.where((TP + FN) == 0, torch.nan, POD)
+        FAR = torch.where((TP + FP) == 0, torch.nan, FAR)
+        ACC = torch.where((TP + FP + TN + FN) == 0, torch.nan, ACC)
+        CSI = torch.where((TP + FP + FN) == 0, torch.nan, CSI)
+
+        return POD, FAR, ACC, CSI
+
+
+def img2tensor(img):
+    if type(img) == str:
+        img = cv2.imread(img, -1)
+    img = torch.from_numpy(img)
+    img = rearrange(img, ' (b c h) w->b c h w', b=1, c=1)
+    return img
+
+
+if __name__ == '__main__':
+    preds = torch.Tensor([[[0, 0, 1],
+                           [0, 1, 1],
+                           [1, 1, 1]
+                           ],
+                          [[0, 1, 1],
+                           [0, 1, 1],
+                           [0, 1, 1]
+                           ]])
+    # target = torch.rand(2, 1, 3, 3)
+    target = torch.Tensor([[[0, 1, 1],
+                            [0, 1, 1],
+                            [0, 1, 1]
+                            ],
+                           [[0, 1, 1],
+                            [0, 1, 1],
+                            [0, 0, 1]
+                            ]])
+    target[1, 0, 2] = torch.nan
+    # m = Metrics()
+    # print(m.calc_rr(preds, target))
+    m = Metrics()
+    AUROC = m.calc_AUROC(preds, target)
+    print(AUROC)
+    print(m.calculate_rain_metrics(preds, target))

jacksung/ai/utils/cmorph.py

@@ -0,0 +1,42 @@
+import netCDF4 as nc
+import numpy as np
+from einops import rearrange, repeat
+from jacksung.utils.data_convert import np2tif, get_transform_from_lonlat_matrices
+
+
+def getNPfromHDF(hdf_path, lock=None, save_file=True, print_log=False):
+    if lock:
+        lock.acquire()
+    ds = nc.Dataset(hdf_path)
+    if lock:
+        lock.release()
+    np_data = np.array(ds['cmorph'][:]).astype(np.float32)
+    lon_array = np.array(ds['lon'][:]).astype(np.float32)
+    lat_array = np.array(ds['lat'][:]).astype(np.float32)
+    lon_dim = len(lon_array)
+    lat_dim = len(lat_array)
+    lon_array = repeat(lon_array, 'w -> w h', h=lat_dim)
+    lat_array = repeat(lat_array, 'h -> w h', w=lon_dim)
+    ds.close()
+    # np_data = rearrange(np_data[0], 'w h->h w')[::-1, :]
+    np_data[np_data < 0] = 0
+    # np_data = np_data[0] + np_data[1]
+    transform, avg_error = get_transform_from_lonlat_matrices(
+        lon_array=lon_array,
+        lat_array=lat_array,
+        gcp_density=20,  # 范围越大，gcp_density建议越大
+        print_log=print_log
+    )
+    if save_file:
+        np2tif(np_data, save_path='np2tif_dir', out_name='CMORPH', dtype='float32', transform=transform)
+    return np_data, transform
+
+
+if __name__ == '__main__':
+    data = getNPfromHDF(rf'C:\Users\ECNU\PycharmProjects\CMORPH_V1.0_ADJ_8km-30min_2022070203.nc')
+    # from datetime import datetime
+    #
+    # da = datetime.utcfromtimestamp(1656730800)
+    # print(da)
+    # da = datetime.utcfromtimestamp(1656732600)
+    # print(da)

jacksung/ai/utils/data_parallelV2.py

@@ -0,0 +1,90 @@
+from torch.nn.parallel import DataParallel
+import torch
+from torch.nn.parallel._functions import Scatter
+from torch.nn.parallel.parallel_apply import parallel_apply
+
+
+def scatter(inputs, target_gpus, chunk_sizes, dim=0):
+    r"""
+    Slices tensors into approximately equal chunks and
+    distributes them across given GPUs. Duplicates
+    references to objects that are not tensors.
+    """
+
+    def scatter_map(obj):
+        if isinstance(obj, torch.Tensor):
+            try:
+                return Scatter.apply(target_gpus, chunk_sizes, dim, obj)
+            except:
+                print('obj', obj.size())
+                print('dim', dim)
+                print('chunk_sizes', chunk_sizes)
+                quit()
+        if isinstance(obj, tuple) and len(obj) > 0:
+            return list(zip(*map(scatter_map, obj)))
+        if isinstance(obj, list) and len(obj) > 0:
+            return list(map(list, zip(*map(scatter_map, obj))))
+        if isinstance(obj, dict) and len(obj) > 0:
+            return list(map(type(obj), zip(*map(scatter_map, obj.items()))))
+        return [obj for targets in target_gpus]
+
+    # After scatter_map is called, a scatter_map cell will exist. This cell
+    # has a reference to the actual function scatter_map, which has references
+    # to a closure that has a reference to the scatter_map cell (because the
+    # fn is recursive). To avoid this reference cycle, we set the function to
+    # None, clearing the cell
+    try:
+        return scatter_map(inputs)
+    finally:
+        scatter_map = None
+
+
+def scatter_kwargs(inputs, kwargs, target_gpus, chunk_sizes, dim=0):
+    r"""Scatter with support for kwargs dictionary"""
+    inputs = scatter(inputs, target_gpus, chunk_sizes, dim) if inputs else []
+    kwargs = scatter(kwargs, target_gpus, chunk_sizes, dim) if kwargs else []
+    if len(inputs) < len(kwargs):
+        inputs.extend([() for _ in range(len(kwargs) - len(inputs))])
+    elif len(kwargs) < len(inputs):
+        kwargs.extend([{} for _ in range(len(inputs) - len(kwargs))])
+    inputs = tuple(inputs)
+    kwargs = tuple(kwargs)
+    return inputs, kwargs
+
+
+class BalancedDataParallel(DataParallel):
+    def __init__(self, gpu0_bsz, *args, **kwargs):
+        self.gpu0_bsz = gpu0_bsz
+        super().__init__(*args, **kwargs)
+
+    def forward(self, *inputs, **kwargs):
+        if not self.device_ids:
+            return self.module(*inputs, **kwargs)
+        gpu_idx_start = 1 if self.gpu0_bsz == 0 else 0
+        inputs, kwargs = self.scatter(inputs, kwargs, self.device_ids[gpu_idx_start:])
+        if len(self.device_ids) == 1:
+            return self.module(*inputs[0], **kwargs[0])
+        replicas = self.replicate(self.module, self.device_ids[:len(inputs) + gpu_idx_start])
+        replicas = replicas[gpu_idx_start:]
+        outputs = self.parallel_apply(replicas, self.device_ids[gpu_idx_start:], inputs, kwargs)
+        return self.gather(outputs, self.output_device)
+
+    def parallel_apply(self, replicas, device_ids, inputs, kwargs):
+        return parallel_apply(replicas, inputs, kwargs, device_ids[:len(inputs)])
+
+    def scatter(self, inputs, kwargs, device_ids):
+        bsz = inputs[0].size(self.dim)
+        num_dev = len(self.device_ids)
+        gpu0_bsz = self.gpu0_bsz
+        bsz_unit = (bsz - gpu0_bsz) // (num_dev - 1)
+        if gpu0_bsz < bsz_unit:
+            chunk_sizes = [gpu0_bsz] + [bsz_unit] * (num_dev - 1)
+            delta = bsz - sum(chunk_sizes)
+            for i in range(delta):
+                chunk_sizes[i + 1] += 1
+            if gpu0_bsz == 0:
+                chunk_sizes = chunk_sizes[1:]
+        else:
+            return super().scatter(inputs, kwargs, device_ids)
+
+        return scatter_kwargs(inputs, kwargs, device_ids, chunk_sizes, dim=self.dim)