kbasic 0.1.16__tar.gz → 0.1.18__tar.gz

{kbasic-0.1.16 → kbasic-0.1.18}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: kbasic
-Version: 0.1.16
+Version: 0.1.18
 Summary: Keyan's basic utility functions.
 Author: Keyan Gootkin
 Author-email: Keyan Gootkin <keyangootkin@gmail.com>

{kbasic-0.1.16 → kbasic-0.1.18}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "kbasic"
-version = "0.1.16"
+version = "0.1.18"
 description = "Keyan's basic utility functions."
 readme = "README.md"
 authors = [

{kbasic-0.1.16 → kbasic-0.1.18}/src/kbasic/Tex.py

@@ -28,7 +28,7 @@ def texfraction(num) -> str:
 # >-|===|>                             Classes                             <|===|-<
 # !==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==
 class Tex(str):
-    def __init__(self, x):
+    def __init__(self, x: str) -> None:
         x = x.strip(' $')
         if x[-1]=='\n': x = x[:-1]
         if x[-2:]=='.0': x = x[:-2]
@@ -42,5 +42,5 @@ class Tex(str):
         self.string = fr"{l2t(x)}"
         self.wrap = "$"+self.string+"$"
 
-    def __repr__(self): return self.string.strip("$")
-    def __str__(self): return self.string
+    def __repr__(self) -> str: return self.string.strip("$")
+    def __str__(self) -> str: return self.string

{kbasic-0.1.16 → kbasic-0.1.18}/src/kbasic/__init__.py

@@ -1,6 +1,7 @@
 from kbasic.environment import *
 from kbasic.audio import *
 from kbasic.array import *
+from kbasic.strings import *
 from kbasic.bar import *
 from kbasic.shell import *
 from kbasic.user_input import *

{kbasic-0.1.16 → kbasic-0.1.18}/src/kbasic/array.py

@@ -1,9 +1,12 @@
 # !==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==
 # >-|===|>                             Imports                             <|===|-<
 # !==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==
+from kbasic.typing import Number, Iterable
 from typing import Callable
-from numpy import ndarray, argmin, any, abs, where, nanmean, nanmin, nanmax, sqrt,\
-                  linspace, nanstd, array, isnan, arange, mgrid, r_, c_, zeros
+from numpy import ndarray, argmin, any, all, absolute, hypot, logspace, log10, \
+                  where, nanmean, nanmin, nanmax, sqrt, linspace, nanstd, array, \
+                  isnan, isfinite, arange, mgrid, r_, c_, zeros
+from numpy.fft import fftshift, fft2
 from scipy.interpolate import RegularGridInterpolator
 
 # !==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==
@@ -15,13 +18,42 @@ from scipy.interpolate import RegularGridInterpolator
 # !==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==
 # >-|===|>                            Functions                            <|===|-<
 # !==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==
-def tile(arr: ndarray) -> ndarray:
-    """take an image and create a 3x3 grid of that image"""
-    return r_[c_[arr, arr, arr], c_[arr, arr, arr], c_[arr, arr, arr]]
-def where_closest(arr:ndarray, x): return int(argmin(abs(arr-x)))
-def where_between(arr:ndarray, low, high): return where((arr>=low)&(arr<=high))
+def tile(arr: Iterable) -> ndarray:
+    """take an array and create a 3x3 grid of copies of that array
+    
+    Args:
+        arr (Iterable): the 2D array you want a tiled version
+
+    Returns:
+        ndarray: a 3x3 grid of coppies of arr
+    """
+    x = array(arr)
+    return r_[c_[x, x, x], c_[x, x, x], c_[x, x, x]]
+def where_closest(arr:ndarray, value: Number): 
+    """find the index where arr is closest to the value of x
+
+    Args:
+        arr (ndarray): the array to search for the index in
+        value (Number): the value you want to get close to
+
+    Returns:
+        _type_: _description_
+    """
+    return int(argmin(absolute(arr-value)))
+def where_between(arr:ndarray, low: Number, high: Number): 
+    """Find the range of indicies where arr is between low and high
+
+    Args:
+        arr (ndarray): the array to search through
+        low (Number): the lower edge of the range (inclusive)
+        high (Number): the upper edge of the range (inclusive)
+
+    Returns:
+        ndarray: an array which is True at indices where low<=arr<=high, and False elsewhere
+    """
+    return where((arr>=low)&(arr<=high))
 def bin_this(
-        x, y, 
+        x: Iterable, y: Iterable, 
         n_bins: int = 50, 
         func: Callable = nanmean
 ) -> tuple[ndarray]:
@@ -50,11 +82,18 @@ def bin_this(
     return x_binned, y_binned, bin_errors
 def nan_clip(*args):
     """
-    take a series of arrays and only return the indicies where ALL members are finite
+    take a series of arrays and only return the indicies where ALL members are not nan
     """
     mask = ~any([isnan(a) for a in args], axis=0)
     nanless_args = tuple([array(a)[mask] for a in args])
     return nanless_args
+def only_finite(*args):
+    """
+    take a series of arrays and only return the indicies where ALL members are finite
+    """
+    mask = all([isfinite(a) for a in args], axis=0)
+    nanless_args = tuple([array(a)[mask] for a in args])
+    return nanless_args
 def interpolate2d(
         data, factor: int, 
         method: str = 'linear',
@@ -81,9 +120,32 @@ def interpolate2d(
     new_grid = mgrid[:Nx:1/factor, :Ny:1/factor].T
     return RegularGridInterpolator(grid, data_repeat, method=method)(new_grid).T
 
-# !==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==
-# >-|===|>                            Decorators                           <|===|-<
-# !==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==
-# !==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==
-# >-|===|>                             Classes                             <|===|-<
-# !==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==
+def kspec(image: ndarray) -> ndarray:
+    """take 2d fft and center k=0 at the center point"""
+    return absolute(fftshift(fft2(image) / (1. * image.shape[0] * image.shape[1])))**2
+
+def kspec1d(image: ndarray, bins: int = 100, return_bin_edges: bool = False) -> tuple[ndarray, ndarray, ndarray]:
+    """Take a 2d image and turn it into a 1d fft
+
+    Args:
+        image (ndarray): the array to take the FFT of 
+        bins (int, optional): number of k bins. Defaults to 100.
+        return_bin_edges (bool, optional): if true return the bin edges, not just the k values. Defaults to False.
+
+    Returns:
+        tuple[ndarray, ndarray, ndarray]: 
+    """
+    k = kspec(image)
+    Ny, Nx = image.shape
+    kmag = hypot(*mgrid[
+                     -Ny // 2: Ny // 2,
+                     -Nx // 2: Nx // 2
+                     ][::-1])
+    kmin = nanmin(kmag[kmag != 0])
+    kmax = nanmax(kmag)
+    kgrid = r_[[0], logspace(log10(3*kmin), log10(kmax), bins)]
+    kx = array([nanmean([kgrid[i], kgrid[i+1]]) for i in range(len(kgrid)-1)])
+    ks = array([nanmean(k[where((kgrid[i] < kmag) & (kmag < kgrid[i+1]))]) for i in range(len(kgrid)-1)])
+    kerr = array([nanstd(in_bin:=k[where((kgrid[i] < kmag) & (kmag < kgrid[i+1]))])/sqrt(len(in_bin)) for i in range(len(kgrid)-1)])
+    if not return_bin_edges: return kx, ks, kerr
+    return kgrid, kx, ks, kerr

{kbasic-0.1.16 → kbasic-0.1.18}/src/kbasic/bar.py

@@ -1,7 +1,7 @@
 # !==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==
 # >-|===|>                             Imports                             <|===|-<
 # !==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==
-from kbasic.typing import Number
+from kbasic.typing import Number, Iterable
 from kbasic.strings import green, yellow, black
 from contextlib import contextmanager
 import inspect
@@ -18,7 +18,7 @@ from tqdm import tqdm
 # !==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==
 def bar(
         x: Number, total: Number, 
-        width: int = 20, border="|", block="▉", color='white'
+        width: int = 20, border: str = "|", block: str = "▉", color: str = 'white'
 ) -> str:
     """create a string representing a progress bar set at 100 * x / total % full.
 
@@ -56,13 +56,13 @@ def redirect_to_tqdm():
         yield
     finally:
         inspect.builtins.print = old_print
-def progress_bar(iterator, **kwargs):
+def progress_bar(iterator: Iterable, **kwargs):
     """tqdm with print redirected to tqdm.write
     """
     with redirect_to_tqdm():
         for x in tqdm(iterator, **kwargs):
             yield x
-def verbose_bar(iterator, verbose, **kwargs):
+def verbose_bar(iterator: Iterable, verbose: bool, **kwargs):
     """just a progress bar if verbose is true.
     """
     return progress_bar(iterator, **kwargs) if verbose else iterator

{kbasic-0.1.16 → kbasic-0.1.18}/src/kbasic/parsing.py

@@ -9,13 +9,12 @@ from glob import glob
 from shutil import copy, move, copytree, rmtree
 from os.path import isdir, isfile, exists, abspath
 from os import system, remove
-from functools import cached_property
 import tomllib
 
 # !==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==
 # >-|===|>                           Definitions                           <|===|-<
 # !==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==
-unreadable_file_types = ['gz', 'tar', 'zip']
+unreadable_file_types: list[str] = ['gz', 'tar', 'zip']
 
 # !==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==
 # >-|===|>                            Functions                            <|===|-<
@@ -27,7 +26,16 @@ def ensure_path(path: str) -> None:
         path (str): the path you want to exist
     """
     system(f"mkdir -p {path}")
-def could_be_path(path: str) -> bool: return isdir('/'.join(path.split('/')[:2])) 
+def could_be_path(path: str) -> bool: 
+    """determine if this is anywhere close to a valid path
+
+    Args:
+        path (str): the path (?) to check
+
+    Returns:
+        bool: True if the first two members of the path are valid else False.
+    """
+    return isdir('/'.join(path.split('/')[:2])) 
 
 # !==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==
 # >-|===|>                             Classes                             <|===|-<
@@ -141,7 +149,7 @@ class Folder:
             return None
         rmtree(self.path)
 
-def parse(path: str | list[str]) -> Folder | File:
+def parse(path: str | list[str]) -> Folder | File | list[Folder | File]:
     """take a path or list of paths and turn them into Folder or File objects as appropriate.
 
     Args:
@@ -151,7 +159,7 @@ def parse(path: str | list[str]) -> Folder | File:
         FileNotFoundError: if you can't match path
 
     Returns:
-        Folder | File: path as a Folder/File.
+        Folder | File | list[Folder|File]: path(s) as a Folder/File.
     """
     match path:
         case str():

{kbasic-0.1.16 → kbasic-0.1.18}/src/kbasic/shell.py

@@ -5,6 +5,7 @@ from kbasic.bar import redirect_to_tqdm
 from kbasic.audio import success
 from kbasic.environment import isAnvil 
 if isAnvil: from kbasic.environment.anvil import anvil_user
+from typing import Any
 from subprocess import check_output, DEVNULL
 from tqdm import tqdm
 from time import sleep
@@ -17,22 +18,38 @@ import numpy as np
 # !==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==
 # >-|===|>                           Definitions                           <|===|-<
 # !==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==
-bad = ['\x1b[31m', '\x1b[34m', '\x1b[m']
-_USERNAME_ = "x-kgootkin" if not isAnvil else anvil_user
+bad: list[str] = ['\x1b[31m', '\x1b[34m', '\x1b[m']
+_USERNAME_: str = "x-kgootkin" if not isAnvil else anvil_user
 # !==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==
 # >-|===|>                            Functions                            <|===|-<
 # !==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==
-def parse_shell_output(output):
-        match output:
-            case str(): return output 
-            case [x]: return x
-            case [x, *_]: 
-                for i in range(len(output)):
-                    for b in bad:
-                        if b in output[i]:
-                            output[i] = output[i].strip(b)
-                return output
-def system(cmd: str):
+def parse_shell_output(output: str | list[str]) -> str | list[str]:
+    """take the output of a shell command and make it nice
+
+    Args:
+        output (str | list[str]): the output of a shell command
+
+    Returns:
+        str | list[str]: a cleaned version of output
+    """
+    match output:
+        case str(): return output 
+        case [x]: return x
+        case [x, *_]: 
+            for i in range(len(output)):
+                for b in bad:
+                    if b in output[i]:
+                        output[i] = output[i].strip(b)
+            return output
+def system(cmd: str) -> str | list[str]:
+    """a version of os.system that actually can return the output
+
+    Args:
+        cmd (str): the command to give the terminal
+
+    Returns:
+        str | list[str]: the output of the command (cleaned of color tags)
+    """
     command = cmd.split(' ') if type(cmd)==str else cmd
     for i in range(len(command)-1):
         if command[i].startswith('"'): 
@@ -43,11 +60,24 @@ def system(cmd: str):
             for j in range(start_quote+1, end_quote+1): del command[j]
     output = parse_shell_output(check_output(command, stderr=DEVNULL).decode().splitlines())
     return output
-def anvil(cmd: str, username=_USERNAME_):
+def anvil(cmd: str, username=_USERNAME_) -> str | list[str]:
+    """Send a shell command to anvil and parse the output
+
+    Args:
+        cmd (str): the command to send to anvil
+        username (str, optional): your anvil username. tries to parse your anvil username by default.
+
+    Returns:
+        str | list[str]: the output of the command (cleaned of color tags)
+    """
     output = parse_shell_output(check_output(['ssh', f'{username}@anvil.rcac.purdue.edu', *cmd.split(' ')], stderr=DEVNULL).decode().splitlines())
     return output
-async def anvil_async(cmd: str): asyncio.to_thread(anvil, cmd)
-def anvil_queue(username=_USERNAME_): return anvil(f"squeue -u {username}")
+async def anvil_async(cmd: str): 
+    """Just an asynchronous version of the anvil command"""
+    asyncio.to_thread(anvil, cmd)
+def anvil_queue(username=_USERNAME_):
+    """Send a squeue command to anvil to check on your runs""" 
+    return anvil(f"squeue -u {username}")
 qs = anvil_queue
 
 # !==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==
@@ -57,7 +87,7 @@ qs = anvil_queue
 # >-|===|>                             Classes                             <|===|-<
 # !==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!== 
 class AnvilJob:
-    def __init__(self, queue_row: str, sep="DISTINCTSEPERATOR"):
+    def __init__(self, queue_row: str, sep="DISTINCTSEPERATOR") -> None:
         self.sep = sep
         [
             jobid, username, account, name, nodes, cpus, time_limit, status, time
@@ -72,7 +102,7 @@ class AnvilJob:
         self.status = str(status)
         self.time = str(time)
 
-    def __repr__(self): return "-"*30 + f"\n{self.name}: {self.status}\n\t{self.time}/{self.time_limit}"
+    def __repr__(self) -> str: return "-"*30 + f"\n{self.name}: {self.status}\n\t{self.time}/{self.time_limit}"
 
     # def update(self, input=True, iter=True):
     #     match input, iter:
@@ -85,11 +115,13 @@ class AnvilJob:
     #             self.input = dHybridRinput(anvil(f"cat /anvil/scratch/{self.username}/sims/{self.name}/input/input"))
     #         case False, True:
     #             self.iter = int(anvil(f"ls /anvil/scratch/{self.username}/sims/{self.name}/Output/Fields/Magnetic/Total/x/")[-1][5:-3])
-def get_anvil_jobs(username=_USERNAME_):
+def get_anvil_jobs(username=_USERNAME_) -> list[AnvilJob]:
+    """Parse the squeue and return AnvilJob objects"""
     q = anvil(f"squeue -u {username}")
     if type(q)==str: return []
     return [AnvilJob(x) for x in q[1:]]
 async def get_anvil_jobs_async(username=_USERNAME_):
+    """Just an asynchronous version of get_anvil_jobs"""
     q = await asyncio.to_thread(anvil, f"squeue -u {username}")
     if type(q)==str: return []
     return [AnvilJob(x) for x in q[1:]]

{kbasic-0.1.16 → kbasic-0.1.18}/src/kbasic/typing.py

@@ -9,7 +9,8 @@
 # >-|===|>                             Imports                             <|===|-<
 # !==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==
 from numpy import ndarray, int8, uint8, int16, uint16, int32, uint32, int64, uint64, float16, float32, float64, longdouble, complex64, complex128, clongdouble
-
+from collections.abc import Callable, Iterable
+from typing import Any
 # !==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==
 # >-|===|>                              Types                              <|===|-<
 # !==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==
@@ -19,20 +20,6 @@ class Number:
         float, float16, float32, float64, longdouble, 
         complex, complex64, complex128, clongdouble
     ]
-class Iterable:
-    types: list = [
-        list, set, dict, tuple, ndarray
-    ]
-
-# !==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==
-# >-|===|>                           Definitions                           <|===|-<
-# !==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==
-# !==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==
-# >-|===|>                            Functions                            <|===|-<
-# !==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==
-# !==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==
-# >-|===|>                            Decorators                           <|===|-<
-# !==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==
-# !==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==
-# >-|===|>                             Classes                             <|===|-<
-# !==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==
+    
+    def __new__(cls):
+        return cls

{kbasic-0.1.16 → kbasic-0.1.18}/src/kbasic/user_input.py

@@ -1,15 +1,20 @@
 # !==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==
 # >-|===|>                             Imports                             <|===|-<
 # !==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==
-# !==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==
-# >-|===|>                              Types                              <|===|-<
-# !==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==
-# !==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==
-# >-|===|>                           Definitions                           <|===|-<
-# !==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==
+from kbasic.typing import Number
+from typing import Any
+
 # !==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==
 # >-|===|>                            Functions                            <|===|-<
 # !==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==
+def parse_user_input(response: str, sep: str = ',') -> Any:
+    match response:
+        case str(x) if x.isnumeric(): 
+            return Number.apply(response)
+        case str(x) if ',' in x: 
+            return tuple(parse_user_input(xi.strip(), sep=sep) for xi in response.split(sep))
+        case _: 
+            return response.lower().strip()
 def yesno(prompt: str):
     """
     prompt the user to either reply yes or no
@@ -33,10 +38,6 @@ def yesno(prompt: str):
                 raise ValueError("need a response with either y or n in it.")
 
         return retry_yesno()
-
-# !==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==
-# >-|===|>                            Decorators                           <|===|-<
-# !==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==
-# !==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==
-# >-|===|>                             Classes                             <|===|-<
-# !==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==!==
+def interactive_set_attribute(obj: Any, attr: str, default_answer: str = "") -> None:
+    res: Any = parse_user_input(input(f"Set a value for {repr(obj)}.{attr}:\n\t"))
+    setattr(obj, attr, res)