├── .gitattributes ├── .github ├── ISSUE_TEMPLATE │ ├── bug_report.md │ └── feature_request.md └── workflows │ ├── run_tests.yaml │ └── test_pre_commit.yaml ├── .gitignore ├── .pre-commit-config.yaml ├── CONTRIBUTING.md ├── LICENSE ├── README.MD ├── assets ├── numojo_logo.png ├── numojo_logo_1.png └── numojo_logo_360x360.png ├── docs ├── changelog.md ├── features.md ├── readme_jp.md ├── readme_kr.md ├── readme_zhs.md ├── readme_zht.md ├── roadmap.md ├── style guide.md ├── style.json └── style_example.mojo ├── mojoproject.toml ├── numojo ├── __init__.mojo ├── core │ ├── __init__.mojo │ ├── _array_funcs.mojo │ ├── _math_funcs.mojo │ ├── complex │ │ ├── __init__.mojo │ │ ├── complex_dtype.mojo │ │ ├── complex_ndarray.mojo │ │ └── complex_simd.mojo │ ├── datatypes.mojo │ ├── flags.mojo │ ├── item.mojo │ ├── matrix.mojo │ ├── ndarray.mojo │ ├── ndshape.mojo │ ├── ndstrides.mojo │ ├── own_data.mojo │ ├── ref_data.mojo │ ├── traits │ │ ├── __init__.mojo │ │ ├── array_like.mojo │ │ ├── backend.mojo │ │ ├── bufferable.mojo │ │ └── indexer_collection_element.mojo │ └── utility.mojo ├── prelude.mojo ├── routines │ ├── __init__.mojo │ ├── bitwise.mojo │ ├── constants.mojo │ ├── creation.mojo │ ├── functional.mojo │ ├── indexing.mojo │ ├── io │ │ ├── __init__.mojo │ │ ├── files.mojo │ │ └── formatting.mojo │ ├── linalg │ │ ├── __init__.mojo │ │ ├── decompositions.mojo │ │ ├── misc.mojo │ │ ├── norms.mojo │ │ ├── products.mojo │ │ └── solving.mojo │ ├── logic │ │ ├── __init__.mojo │ │ ├── comparison.mojo │ │ ├── contents.mojo │ │ └── truth.mojo │ ├── manipulation.mojo │ ├── math │ │ ├── __init__.mojo │ │ ├── arithmetic.mojo │ │ ├── differences.mojo │ │ ├── exponents.mojo │ │ ├── extrema.mojo │ │ ├── floating.mojo │ │ ├── hyper.mojo │ │ ├── misc.mojo │ │ ├── products.mojo │ │ ├── rounding.mojo │ │ ├── sums.mojo │ │ └── trig.mojo │ ├── random.mojo │ ├── searching.mojo │ ├── sorting.mojo │ └── statistics │ │ ├── __init__.mojo │ │ └── averages.mojo └── science │ ├── __init__.mojo │ ├── interpolate.mojo │ └── signal.mojo └── tests ├── README.md ├── core ├── test_array_indexing_and_slicing.mojo ├── test_array_methods.mojo ├── test_bool_masks.mojo ├── test_complexArray.mojo ├── test_complexSIMD.mojo ├── test_matrix.mojo └── test_shape_strides_item.mojo ├── routines ├── test_creation.mojo ├── test_functional.mojo ├── test_indexing.mojo ├── test_linalg.mojo ├── test_manipulation.mojo ├── test_math.mojo ├── test_random.mojo ├── test_sorting.mojo └── test_statistics.mojo ├── science └── test_signal.mojo └── utils_for_test.mojo /.gitattributes: -------------------------------------------------------------------------------- 1 | docs/readthedocs/docs.json -diff -merge -text 2 | docs/readthedocs/docs.json linguist-generated=true 3 | docs/readthedocs/docs/autodocs/*.md -diff -merge -text 4 | docs/readthedocs/docs/autodocs/*.md linguist-generated=true 5 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/bug_report.md: -------------------------------------------------------------------------------- 1 | --- 2 | name: Bug report 3 | about: Create a report to help us improve 4 | title: '' 5 | labels: '' 6 | assignees: '' 7 | 8 | --- 9 | 10 | **Describe the bug** 11 | A clear and concise description of what the bug is. 12 | 13 | **To Reproduce** 14 | Steps to reproduce the behavior: 15 | 1. Go to '...' 16 | 2. Click on '....' 17 | 3. Scroll down to '....' 18 | 4. See error 19 | 20 | **Expected behavior** 21 | A clear and concise description of what you expected to happen. 22 | 23 | **Screenshots** 24 | If applicable, add screenshots to help explain your problem. 25 | 26 | **Desktop (please complete the following information):** 27 | - OS: [e.g. iOS] 28 | - Browser [e.g. chrome, safari] 29 | - Version [e.g. 22] 30 | 31 | 32 | **Additional context** 33 | Add any other context about the problem here. 34 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/feature_request.md: -------------------------------------------------------------------------------- 1 | --- 2 | name: Feature request 3 | about: Suggest an idea for this project 4 | title: '' 5 | labels: '' 6 | assignees: '' 7 | 8 | --- 9 | 10 | **Is your feature request related to a problem? Please describe.** 11 | A clear and concise description of what the problem is. Ex. I'm always frustrated when [...] 12 | 13 | **Describe the solution you'd like** 14 | A clear and concise description of what you want to happen. 15 | 16 | **Describe alternatives you've considered** 17 | A clear and concise description of any alternative solutions or features you've considered. 18 | 19 | **Additional context** 20 | Add any other context or screenshots about the feature request here. 21 | -------------------------------------------------------------------------------- /.github/workflows/run_tests.yaml: -------------------------------------------------------------------------------- 1 | name: NuMojo Unit Tests 2 | on: 3 | pull_request: 4 | workflow_dispatch: 5 | 6 | permissions: 7 | contents: read 8 | pull-requests: read 9 | 10 | jobs: 11 | testing-numojo: 12 | name: with ${{ matrix.os }} 13 | strategy: 14 | fail-fast: false 15 | matrix: 16 | # os: ["ubuntu-22.04", "macos-14"] 17 | os: ["ubuntu-22.04"] 18 | 19 | runs-on: ${{ matrix.os }} 20 | timeout-minutes: 30 21 | 22 | defaults: 23 | run: 24 | shell: bash 25 | env: 26 | DEBIAN_FRONTEND: noninteractive 27 | 28 | steps: 29 | - name: Checkout repo 30 | uses: actions/checkout@v4 31 | 32 | - name: Install magic 33 | run: | 34 | curl -ssL https://magic.modular.com/deb181c4-455c-4abe-a263-afcff49ccf67 | bash 35 | 36 | - name: Add path 37 | run: | 38 | echo "MODULAR_HOME=$HOME/.modular" >> $GITHUB_ENV 39 | echo "$HOME/.modular/bin" >> $GITHUB_PATH 40 | echo "$HOME/.modular/pkg/packages.modular.com_mojo/bin" >> $GITHUB_PATH 41 | 42 | - name: Activate virtualenv 43 | run: | 44 | python3 -m venv $HOME/venv/ 45 | . $HOME/venv/bin/activate 46 | echo PATH=$PATH >> $GITHUB_ENV 47 | 48 | - name: Install packages 49 | run: | 50 | pip install "numpy" 51 | 52 | - name: Run tests 53 | run: | 54 | magic install 55 | magic run mojo test tests -I . 56 | -------------------------------------------------------------------------------- /.github/workflows/test_pre_commit.yaml: -------------------------------------------------------------------------------- 1 | name: Run pre-commit 2 | on: 3 | # Run pre-commit on pull requests 4 | pull_request: 5 | # Add a workflow_dispatch event to run pre-commit manually 6 | workflow_dispatch: 7 | 8 | permissions: 9 | contents: read 10 | pull-requests: read 11 | 12 | jobs: 13 | lint: 14 | runs-on: "ubuntu-22.04" 15 | timeout-minutes: 30 16 | 17 | defaults: 18 | run: 19 | shell: bash 20 | env: 21 | DEBIAN_FRONTEND: noninteractive 22 | 23 | steps: 24 | - name: Checkout repo 25 | uses: actions/checkout@v4 26 | 27 | - name: Install magic 28 | run: | 29 | curl -ssL https://magic.modular.com/deb181c4-455c-4abe-a263-afcff49ccf67 | bash 30 | 31 | - name: Add path 32 | run: | 33 | echo "MODULAR_HOME=$HOME/.modular" >> $GITHUB_ENV 34 | echo "$HOME/.modular/bin" >> $GITHUB_PATH 35 | echo "$HOME/.modular/pkg/packages.modular.com_mojo/bin" >> $GITHUB_PATH 36 | 37 | - name: Activate virtualenv 38 | run: | 39 | python3 -m venv $HOME/venv/ 40 | . $HOME/venv/bin/activate 41 | echo PATH=$PATH >> $GITHUB_ENV 42 | 43 | - name: Install pre-commit 44 | run: | 45 | pip install pre-commit 46 | pre-commit install 47 | 48 | - name: Run pre-commit 49 | run: | 50 | magic install 51 | pre-commit run --all-files -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Build directory 2 | /build 3 | /dist 4 | /local 5 | *.egg-info 6 | 7 | # VSCode files 8 | /.vscode 9 | 10 | # MacOs Desktop Service file 11 | *.DS_Store 12 | 13 | # magic environments 14 | .magic 15 | magic.lock 16 | 17 | # pixi environments 18 | .pixi 19 | /venv 20 | 21 | 22 | # Irrelevant files 23 | *.html 24 | *.css 25 | *.py 26 | *.pyc 27 | 28 | # Miscellaneous files 29 | mojo 30 | numojo.mojopkg 31 | /bench.mojo 32 | /test*.mojo 33 | /test*.ipynb 34 | /tempCodeRunnerFile.mojo 35 | 36 | # Auto docs 37 | docs/readthedocs/docs.json 38 | docs.json -------------------------------------------------------------------------------- /.pre-commit-config.yaml: -------------------------------------------------------------------------------- 1 | repos: 2 | - repo: local 3 | hooks: 4 | - id: mojo-format 5 | name: mojo-format 6 | entry: magic run mojo format 7 | language: system 8 | files: '\.(mojo|🔥|py)$' 9 | stages: [pre-commit] 10 | # - id: autodoc 11 | # name: mautodoc 12 | # entry: magic run doc_pages 13 | # language: system 14 | # files: '\.(mojo|🔥|py)$' 15 | # stages: [pre-commit] 16 | -------------------------------------------------------------------------------- /CONTRIBUTING.md: -------------------------------------------------------------------------------- 1 | # Contributing to NuMojo 2 | 3 | Thank you for your interest in contributing to NuMojo! We appreciate your efforts to make this project better. Below are some guidelines to help you contribute effectively. 4 | 5 | ## Style Guide 6 | 7 | Please follow the Mojo standard library style guide for all contributions. Consistency in code style helps maintain readability and ease of collaboration. Key points include: 8 | 9 | - Use clear, descriptive names for variables and functions. 10 | - Write concise, well-documented code. 11 | - Adhere to formatting conventions for indentation, spacing, and line breaks. 12 | 13 | Additionally refer to `style guide.md` for docstring and naming conventions. 14 | 15 | ## Pull Requests 16 | 17 | When submitting pull requests: 18 | 19 | - Ensure they are small but complete. A pull request should introduce at least a minimally viable version of the feature you are adding. 20 | - Include tests for your contributions where applicable. 21 | - Provide a clear and concise description of what your pull request accomplishes. 22 | 23 | ## Just Do It 24 | 25 | If you have an idea or want to work on something, go ahead and do it! You don’t need to ask for permission before starting. In fact, we prefer if you avoid “licking the cookie” by claiming tasks without following through. We would rather recieve 5 different ways of accomplishing something and then choose the best one or combine then than not recieve any feature at all. 26 | 27 | ## Directory Structure 28 | 29 | Organize your additions into the appropriate submodule or file, if one does not exist feel free to make it and we can figure out where it goes during the pull request checks. This helps keep the project structured and maintainable. For example: 30 | 31 | - If you’re adding a statistics function, place it in the `stats` submodule. 32 | - If a stats module does not yet exist put the code in a directory called stats in a file with a name that describes the sub disipline of statistics the code enables, along with a `__init__.mojo` 33 | - For a kernel density estimation function, add it to the `kde.mojo` file within the `stats` directory. 34 | 35 | Following this structure ensures that similar functionalities are grouped together, making the codebase easier to navigate. 36 | 37 | ## Contribution Process 38 | 39 | 1. **Fork the Repository**: Create a personal fork of the repository on GitHub. 40 | 2. **Clone Your Fork**: Clone your forked repository to your local machine. 41 | 42 | ```sh 43 | git clone https://github.com/your-username/numojo.git 44 | ``` 45 | 46 | 3. **Create a Branch**: Create a new branch for your feature or bugfix. 47 | 48 | ```sh 49 | git checkout -b feature-name 50 | ``` 51 | 52 | 4. **Make Your Changes**: Implement your changes in your branch. 53 | 5. **Run Tests**: NuMojo now uses the `Magic` package manager by Modular. To ensure that all unit tests pass, the NuMojo module packages correctly, and the .mojo files are properly formatted, run the following command: 54 | ```sh 55 | magic run final 56 | ``` 57 | 6. **Commit Your Changes**: Commit your changes with a clear and descriptive commit message. 58 | 59 | ```sh 60 | git commit -m "Add feature XYZ" 61 | ``` 62 | 63 | 7. **Push Your Changes**: Push your branch to your fork on GitHub. 64 | 65 | ```sh 66 | git push origin feature-name 67 | ``` 68 | 69 | 8. **Submit a Pull Request**: Open a pull request to the `main` branch of the original repository. 70 | -------------------------------------------------------------------------------- /assets/numojo_logo.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Mojo-Numerics-and-Algorithms-group/NuMojo/f0dea371bfd15df65e917e013e11d6177bfc975a/assets/numojo_logo.png -------------------------------------------------------------------------------- /assets/numojo_logo_1.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Mojo-Numerics-and-Algorithms-group/NuMojo/f0dea371bfd15df65e917e013e11d6177bfc975a/assets/numojo_logo_1.png -------------------------------------------------------------------------------- /assets/numojo_logo_360x360.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Mojo-Numerics-and-Algorithms-group/NuMojo/f0dea371bfd15df65e917e013e11d6177bfc975a/assets/numojo_logo_360x360.png -------------------------------------------------------------------------------- /docs/features.md: -------------------------------------------------------------------------------- 1 | # Available functions and objects by topics 2 | 3 | Available functions and objects by topics (also see imports within `__init__` file [here](../numojo/__init__.mojo)): 4 | 5 | - Array creation routines (`numojo.routines.creation`) 6 | - zeros, eye, identity, ones, fill 7 | - arange, linspace, logspace, geomspace 8 | - Array manipulation routines (`numojo.routines.manipulation`) 9 | - Bit-wise operations (`numojo.routines.bitwise`) 10 | - Constants (`numojo.routines.constants`) 11 | - Indexing routines (`numojo.routines.indexing`) 12 | - Input and output (`numojo.routines.io`) 13 | - Text files (`numojo.routines.files`) 14 | - Text formatting options (`numojo.routines.formatting`) 15 | - Linear algebra (`numojo.routines.linalg`) 16 | - Decompositions (`numojo.routines.decompositions`) 17 | - Products of matrices and vectors (`numojo.routines.products`) 18 | - Solving (`numojo.routines.solving`) 19 | - Logic functions (`numojo.routines.logic`) 20 | - Comparison (`numojo.routines.comparison`) 21 | - Array contents (`numojo.routines.contents`) 22 | - Truth value testing (`numojo.routines.truth`) 23 | - Mathematical functions (`numojo.routines.math`) 24 | - Sums (`numojo.routines.sums`) 25 | - `sum()` 26 | - Products (`numojo.routines.products`) 27 | - `prod()` 28 | - Differences (`numojo.routines.differences`) 29 | - Extrema (`numojo.routines.math.extrema`) 30 | - `max()`, `min()` 31 | - Trigonometry: acos, asin, atan, cos, sin, tan, atan2, hypot 32 | - Hyperbolic: acosh, asinh, atanh, cosh, sinh, tanh 33 | - Floating: copysign 34 | - exponential and logarithmic: log10, log1p, exp2, exp, log, log2 35 | - Miscellaneous: cbrt, expm1, rsqrt, scalb 36 | - Arithmetic operations (`numojo.routines.arithmetic`) 37 | - pow, mod, mul, sub, add, div 38 | - Exponents and logarithms (`numojo.routines.exponents`) 39 | - Extrema finding (`numojo.routines.extrema`) 40 | - Floating point routines (`numojo.routines.floating`) 41 | - Hyperbolic functions (`numojo.routines.hyper`) 42 | - Indexing (`numojo.routines.indexing`) 43 | - Miscellaneous (`numojo.routines.misc`) 44 | - Rounding (`numojo.routines.rounding`) 45 | - Trigonometric functions (`numojo.routines.trig`) 46 | - Random sampling (`numojo.routines.random`) 47 | - Sorting, searching, and counting (`numojo.routines.sorting`, `numojo.routines.searching`) 48 | - Statistics (`numojo.routines.statistics`) 49 | - Averages and variances (`numojo.routines.averages`) 50 | - `mean()`, `mode()`, `median()` 51 | - `variance()`, `std()` 52 | 53 | To-be-implemented functions and objects by topics: 54 | 55 | - Mathematical functions: abs, floor, ceil, trunc, round, roundeven, round_half_down, round_half_up, reciprocal, nextafter, remainder 56 | -------------------------------------------------------------------------------- /docs/readme_jp.md: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 |
6 | 7 | Logo 8 | 9 | 10 |

NuMojo

11 | 12 |

13 | NuMojoは、PythonのNumPyやSciPyに似たMojo🔥で数値計算を行うためのライブラリです. 14 |
15 | 16 | ドキュメントを読む» 17 |
18 | 19 | Discord チャンネルに参加する» 20 |
21 | 22 | 27 |

28 |
29 | 30 | 52 | 53 | ## プロジェクトについて 54 | 55 | ### NuMojoとは 56 | 57 | NuMojoは、PythonのNumPy、SciPyとScikit に存在する幅広い数値機能を取り込むことを目的としています。 58 | 59 | ベクトル化、並列化、GPUアクセラレーション（利用可能になった場合）など、Mojoの機能を最大限に活用することを試みています。現在、NuMojoは、配列入力で動作するようにスタンダードライブラリの数学関数を（ほとんど）拡張しています。 60 | 61 | NuMojoは、MLのバックとフォワード伝搬システムの負荷なしに高速な計算を必要とする他のMojoパッケージのためのビルディングブロックになることを意図している 62 | 63 | 注意：NuMojoは機械学習ライブラリではなく、コアライブラリに機械学習アルゴリズムが含まれることはありません。 64 | 65 | ## 目標 66 | 67 | 詳細なロードマップについては、[roadmap.md](roadmap.md)(英語)ファイルを参照してください。 68 | 69 | 私たちの主な目標は、Mojoに高速で包括的な数値計算ライブラリを実装することです。以下はNuMojoの長期目標です、 70 | 71 | ###　長期目標 72 | 73 | * 線形代数 74 | * ネイティブの n 次元配列 75 | * ベクトル化、並列化された数学演算 76 | * 配列操作 - vstack、スライス、連結など 77 | * 微積分 78 | * 積分と微分など 79 | * オプティマイザ 80 | * 関数近似 81 | * 並べ替え 82 | 83 | ## 使い方 84 | 85 | 以下にコード例を示します、 86 | 87 | ```mojo 88 | import numojo as nm 89 | 90 | fn main() raises: 91 | # ランダムな float64 値を使用して 2 つの 1000x1000 行列を生成する。 92 | var A = nm.NDArray[nm.f64](shape=List[Int](1000,1000), random=True) 93 | var B = nm.NDArray[nm.f64](1000,1000, random=True) 94 | 95 | # A*B 96 | print(nm.linalg.matmul_parallelized(A, B)) 97 | ``` 98 | 99 | 利用可能なすべての機能は[ここ](features.md)で見つけてください 100 | 101 | ## インストール方法 102 | 103 | NuMojoパッケージをインストールして利用するには２つの方法があります。 104 | 105 | ### パッケージのビルド方法 106 | 107 | このアプローチでは、スタンドアロンパッケージファイル `mojopkg` をビルドする。 108 | 109 | 1. リポジトリをクローンする。 110 | 2. `mojo pacakge numojo` を使用してパッケージをビルドする。 111 | 3. numojo.mojopkg をあなたのコードを含むディレクトリに移動する。 112 | 113 | ### コンパイラとLSPにNuMojoのパスを含める。 114 | 115 | この方法では、パッケージファイルを作成する必要はありません。コードをコンパイルするときに、以下のコマンドでNuMojoリポジトリのパスをインクルードできます： 116 | 117 | ```console 118 | mojo run -I "../NuMojo" example.mojo 119 | ``` 120 | 121 | これは、コードをテストするときにNuMojoソースファイルを編集できるので、より柔軟です。 122 | 123 | VSCode LSPがインポートされた `numojo` パッケージを解決できるようにするには、次のようにします： 124 | 125 | 1. VSCodeの環境設定ページを開きます。 126 | 2. Mojo ' Lsp: Include Dirs` に移動します。 127 | 3. add item` をクリックし、Numojo リポジトリがあるパスを追加します。例えば `/Users/Name/Programs/NuMojo` です。 128 | 。 129 | 4. Mojo LSPサーバーを再起動します。 130 | 131 | これでVSCodeがNumojoパッケージの関数ヒントを表示できるようになります！ 132 | 133 | ## 貢献 134 | 135 | どのような貢献でも大歓迎です**。コントリビュートに関する詳細やガイドラインは、[こちら](CONTRIBUTING.md)を参照してください。 136 | 137 | ## 警告 138 | 139 | このライブラリはまだ非常に未完成であり、いつでも変更される可能性があります。 140 | 141 | ## ライセンス 142 | 143 | LLVM例外を含むApache 2.0ライセンスの下で配布されています。詳細は[LICENSE](https://github.com/Mojo-Numerics-and-Algorithms-group/NuMojo/blob/main/LICENSE)とLLVM [License](https://llvm.org/LICENSE.txt)を参照してください。 144 | 145 | ## 謝辞 146 | 147 | * Modular](https://github.com/modularml)によって作成されたネイティブの[Mojo](https://github.com/modularml/mojo)で構築されています。 -------------------------------------------------------------------------------- /docs/readme_kr.md: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 |
6 | 7 | Logo 8 | 9 | 10 |

NuMojo

11 | 12 |

13 | NuMojo는 Python의 NumPy, SciPy와 유사한 Mojo🔥의 수치 컴퓨팅용 라이브러리입니다. 14 |
15 | 16 | Explore the docs» 17 |
18 | Check out our Discord» 19 |
20 | 21 | 26 |

27 |
28 | 29 | -------------------------------------------------------------------------------- /docs/readme_zhs.md: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 |
6 | 7 | Logo 8 | 9 | 10 |

NuMojo

11 | 12 |

13 | NuMojo 是为 Mojo 🔥 设计的多维数组运算库，类似 Python 中的 NumPy, SciPy。 14 |
15 |

16 | 阅读文档»     17 | 更新日志»     18 | 加入 Discord 讨论频道» 19 |
20 |
21 |
22 | 中文·繁»   23 | 日本語» 24 | English»   25 |
26 |

27 |
28 | 29 | ## 关于本项目 30 | 31 | NuMojo，旨在为 Mojo 语言生态系统提供数值计算和多维数组运算功能，类似于 NumPy、SciPy 和 Scikit 等数学库在 Python 语言生态系统中所扮演的角色。 32 | 33 | NuMojo 充分利用 Mojo 的潜力，包括向量化、并行化和 GPU 加速。目前，NuMojo 扩展了大部分 Mojo 标准库中的数学函数，用以处理多维数组。 34 | 35 | NuMojo 也可为其他需要高速数值计算、多维数组运算等功能的 Mojo 第三方库提供基础类型和函数。 36 | 37 | 注意：NuMojo 不是一个机器学习库，它永远不会在核心库中包含机器学习算法。 38 | 39 | ## 目标及路线图 40 | 41 | 有关本项目详细的路线图，请参阅 [roadmap.md](../roadmap.md) 文件（英文）。 42 | 43 | 我们的核心目标，是使用 Mojo 实现一个快速、全面的数值计算库。以下是部分长期目标： 44 | 45 | - 原生 N 维数组类型 46 | - 向量化、并行化的数值运算 47 | - 线性代数 48 | - 数组操作：叠加、切片、拼合等 49 | - 微积分 50 | - 优化 51 | - 函数逼近和估值 52 | - 排序 53 | 54 | ## 使用方法 55 | 56 | 以下为部分代码实例： 57 | 58 | ```mojo 59 | import numojo as nm 60 | from numojo.prelude import * 61 | 62 | fn main() raises: 63 | # 生成两个 1000x1000 矩阵，数值随机且为 64 位浮点数 64 | var A = nm.random.randn[f64](shape=List[Int](1000, 1000)) 65 | var B = nm.random.randn[f64](shape=List[Int](1000, 1000)) 66 | 67 | # 根据字符串生成 3x2 矩阵，数据类型为 32 位浮点数 68 | var X = nm.fromstring[f32]("[[1.1, -0.32, 1], [0.1, -3, 2.124]]") 69 | 70 | # 打印矩阵 71 | print(A) 72 | 73 | # 矩阵相乘 74 | var C = A @ B 75 | 76 | # 矩阵求逆 77 | var I = nm.inv(A) 78 | 79 | # 矩阵切片 80 | var A_slice = A[1:3, 4:19] 81 | 82 | # 提取矩阵元素 83 | var A_item = A.item(291, 141) 84 | ``` 85 | 86 | 请在 [此文档](../features.md) 中查询所有可用的函数。 87 | 88 | ## 安装方法 89 | 90 | Numojo 库可通过两种方法安装并使用。 91 | 92 | ### 构建文件包 93 | 94 | 这种方法会构建一个独立文件包 `numojo.mojopkg`。步骤为： 95 | 96 | 1. 克隆本仓库。 97 | 1. 在控制台使用 `mojo package numojo` 命令构建文件包。 98 | 1. 将 `numojo.mojopkg` 移动到包含代码的目录中，即可使用。 99 | 100 | ### 将 NuMojo 路径添加至编译器和 LSP 101 | 102 | 这种方法不需要生成文件包，仅需在编译时，通过以下命令指明 `Numojo` 的文件路径： 103 | 104 | ```console 105 | mojo run -I "../NuMojo" example.mojo 106 | ``` 107 | 108 | 这种方法自由度更高，允许你在调试你代码的过程中修改 `NuMojo` 源码。它适合想要为本库贡献代码的用户。 109 | 110 | 为了使 VSCode 的 LSP （语言服务引擎）解析 `numojo` 库，你可以： 111 | 112 | 1. 进入 VSCode 的 Preference 页面（偏好设置）。 113 | 1. 选择 `Mojo › Lsp: Include Dirs` 114 | 1. 点击 `add item`，并添加 NuMojo 源码所在的路径，例如 `/Users/Name/Programs/NuMojo` 115 | 1. 重启 Mojo LSP server 116 | 117 | 如此，VSCode 便可以提供 NuMojo 包的函数提示。 118 | -------------------------------------------------------------------------------- /docs/readme_zht.md: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 |
6 | 7 | Logo 8 | 9 | 10 |

NuMojo

11 | 12 |

13 | NuMojo 是爲 Mojo 🔥 設計的多維數組運算庫，類似 Python 中的 NumPy, SciPy。 14 |
15 |

16 | 閲讀文檔»     17 | 更新日誌»     18 | 加入 Discord 討論頻道» 19 |
20 |
21 |
22 | 中文·简» 23 | 日本語» 24 | English»   25 |
26 |

27 |
28 | 29 | ## 關於本項目 30 | 31 | NuMojo，旨在為 Mojo 語言生態系統提供數值計算和多維數組運算功能，類似於 NumPy、SciPy 和 Scikit 等數學庫在 Python 語言生態系統中所扮演的角色。 32 | 33 | NuMojo 充分利用 Mojo 的潛力，包括向量化、並行化和 GPU 加速。目前，NuMojo 擴展了大部分 Mojo 標準庫中的數學函數，用以處理多維數組。 34 | 35 | NuMojo 也可為其他需要高速數值計算、多維數組運算等功能的 Mojo 第三方庫提供基礎類型和函數。 36 | 37 | 注意：NuMojo 不是一個機器學習庫，它永遠不會在核心庫中包含機器學習算法。 38 | 39 | ## 目标及路线图 40 | 41 | 有關本項目詳細的路線圖，請參閱 [roadmap.md](../roadmap.md) 文件（英文）。 42 | 43 | 我們的核心目標，是使用 Mojo 實現一個快速、全面的數值計算庫。以下是部分長期目標： 44 | 45 | - 原生 N 維數組類型 46 | - 向量化、並行化的數值運算 47 | - 線性代數 48 | - 數組操作：疊加、切片、拼合等 49 | - 微積分 50 | - 優化 51 | - 函數逼近和估值 52 | - 排序 53 | 54 | ## 使用方法 55 | 56 | 以下爲部分代碼實例： 57 | 58 | ```mojo 59 | import numojo as nm 60 | 61 | fn main() raises: 62 | # 生成兩個 1000x1000 矩陣，數值隨機且爲 64 位浮點數 63 | var A = nm.random.randn[f64](shape=List[Int](1000, 1000)) 64 | var B = nm.random.randn[f64](shape=List[Int](1000, 1000)) 65 | 66 | # 根據字符串生成 3x2 矩陣，数據類型爲 32 位浮點數 67 | var X = nm.fromstring[f32]("[[1.1, -0.32, 1], [0.1, -3, 2.124]]") 68 | 69 | # 打印矩陣 70 | print(A) 71 | 72 | # 矩陣相乘 73 | var C = A @ B 74 | 75 | # 矩陣求逆 76 | var I = nm.inv(A) 77 | 78 | # 矩陣切片 79 | var A_slice = A[1:3, 4:19] 80 | 81 | # 提取矩陣元素 82 | var A_item = A.at(291, 141) 83 | ``` 84 | 85 | 請在 [此文檔](../features.md) 中查詢所有可用的函數。 86 | 87 | ## 安裝方法 88 | 89 | Numojo 库可通过两種方法安裝並使用。 90 | 91 | ### 構建文件包 92 | 93 | 这種方法會構建一個獨立文件包 `numojo.mojopkg`。步骤爲： 94 | 95 | 1. 克隆本倉庫。 96 | 1. 在控制臺使用 `mojo package numojo` 命令構建文件包。 97 | 1. 將 `numojo.mojopkg` 移動到包含代碼的目録中，即可使用。 98 | 99 | ### 將 NuMojo 路徑添加至編譯器和 LSP 100 | 101 | 這種方法不需要生成文件包，僅需在編譯時，通這以下命令指明 `Numojo` 的文件路徑： 102 | 103 | ```console 104 | mojo run -I "../NuMojo" example.mojo 105 | ``` 106 | 107 | 這種方法自由度更高，允許你在調試你代碼的過程中修改 `NuMojo` 源碼。它適合想要爲本庫贡獻代碼的用户。 108 | 109 | 爲了使 VSCode 的 LSP （語言服務引擎）解析 `numojo` 庫，你可以： 110 | 111 | 1. 進入 VSCode 的 Preference 頁面（偏好設置）。 112 | 1. 選擇 `Mojo › Lsp: Include Dirs` 113 | 1. 點擊 `add item`，並添加 NuMojo 源碼所在的路徑，例如 `/Users/Name/Programs/NuMojo` 114 | 1. 重啓 Mojo LSP server 115 | 116 | 如此，VSCode 便可以提供 NuMojo 包的函數提示。 117 | -------------------------------------------------------------------------------- /docs/roadmap.md: -------------------------------------------------------------------------------- 1 | # Roadmap 2 | 3 | Our primary objective is to develop a robust and comprehensive library for numerical and scientific computing. To facilitate a smooth transition from Python, we will primarily adhere to the conventions of `numpy` and `scipy`, while allowing deviations when they: 4 | 5 | - Align with the design philosophy of the Mojo programming language. 6 | - Significantly enhance performance, such as through the use of generics. 7 | - Improve consistency in conventions, such as naming parameters and arguments. 8 | 9 | NuMojo is currently in its early development stages. At this point, our focus is to ***stabilize the API*** and ***implement the majority of functionalities***. If you notice any missing features, please consider developing them and submitting a pull request or creating a feature request. 10 | 11 | ## Core Tasks 12 | 13 | - Implement the n-dimensional array type and support SIMD-compatible standard library math functions[^stdlib]. 14 | - Develop `numpy`-like functions for mathematical, statistical, linear algebra, searching, sorting, etc. 15 | - Create `scipy`-like functions for scientific purposes, such as optimizers, function approximators, and FFT. 16 | 17 | ### N-dimensional Arrays 18 | 19 | We have implemented basic functions and methods for the N-dimensional array `NDArray` (and also `ComplexNDArray` and `Matrix`). We are working on incorporating additional essential features similar to those in `numpy`. 20 | 21 | Currently, operations on an array return a copy. When the Mojo programming language supports parameterized traits, some operations (e.g., slicing and transpose) will return a view of the array. This will avoid excessive copying of data, increase memory reuse, and potentially enhance performance. 22 | 23 | In the future, when the Mojo programming language supports GPU functionality as it currently does with SIMD, NuMojo will also provide an option to use the GPU for calculations. 24 | 25 | ### Implement Basic Numeric Functions 26 | 27 | We are currently working on implementing basic numeric functions into NuMojo. The scope is similar to `numpy`. Functions on [this page](https://numpy.org/doc/stable/reference/routines.html) will be considered for gradual implementation in NuMojo. 28 | 29 | ### Implement Advanced Functions 30 | 31 | We also aim to implement advanced functions into NuMojo. The scope is similar to `scipy`. 32 | 33 | ## Internal Organization of Objects and Functions 34 | 35 | NuMojo organizes modules internally according to the following structure[^numpy]: 36 | 37 | 1. A `routines` folder is created under `/numojo`. Functions covered by [this page](https://numpy.org/doc/stable/reference/routines.html) will be considered for implementation in this folder. 38 | 2. Sub-folders[^single] will be created under `/routines` for each topic [on this page](https://numpy.org/doc/stable/reference/routines.html). Examples include: 39 | - `/creation` (Array creation routines) 40 | - `/logic` (Logic functions) 41 | - `/mathematics` (Mathematical functions) 42 | - ... 43 | 3. In each sub-folder, functions are grouped by topics into single Mojo files. For example, in the `/mathematics` folder, the following files will be created [(as classified by NumPy on this page)](https://numpy.org/doc/stable/reference/routines.math.html): 44 | - `trig.mojo` (Trigonometric functions) 45 | - `hyperbolic.mojo` (Hyperbolic functions) 46 | - `exp_log.mojo` (Exponents and logarithms) 47 | - `other.mojo` (Other special functions) 48 | - `arithmetic.mojo` (Arithmetic operations) 49 | - ... 50 | 4. In each file, functions are sorted alphabetically. 51 | 5. The `__init__.mojo` files of parent folders import functions from their child modules explicitly, avoiding `import *` to prevent polluting the namespace. 52 | 53 | Additionally, a `science` folder is created under `/numojo`. It is similar to the `routines` folder but contains sub-packages for features present in `scipy`[^science]. For example: 54 | 55 | Users can access the functions either directly at the top level or via sub-packages. 56 | 57 | 1. Most common functions can be called from the top level, e.g., `numojo.sort()`. 58 | 2. Advanced features (e.g., those listed as sub-packages in `numpy` or `scipy`) need to be called via their own namespaces. For example: 59 | - Random array generators, e.g., `numojo.random.randint()`. 60 | - Linear algebra, e.g., `numojo.linalg.solve()`. 61 | - FFT, e.g., `numojo.fft()`. 62 | - Ordinary differential equations. 63 | - Optimizers, e.g., `numojo.optimize`. 64 | 65 | [^stdlib]: Standard library functions that are SIMD-compatible. 66 | [^numpy]: The structure is inspired by the organization of functions in NumPy. 67 | [^single]: If a topic has only a few functions, they can be grouped into a single Mojo file instead of creating a separate folder. 68 | [^science]: This folder can be further integrated into the `routines` folder if necessary. 69 | -------------------------------------------------------------------------------- /docs/style guide.md: -------------------------------------------------------------------------------- 1 | # Numojo Style Guide 2 | 3 | In the interest of keeping our code clean and consistent, and enabling some automation for documentation the following simple standards will be required for new commits. 4 | 5 | ## File Level 6 | All files must begin with a triple quoted docstring describing the functionality created by the file. It should be a single sentence with the first letter capitalized and ending with a period. 7 | ```python 8 | """ 9 | Document docstring describing what it does, if it is in an init file it will be the docstring for the module. 10 | """ 11 | ``` 12 | All aliases and file-level variable definitions must have a docstring that describes what they are placed below the declaration. 13 | ```python 14 | alias Example = Int 15 | """ Aliases can be explained with docstrings and should if they exist in the global scope.""" 16 | ``` 17 | Aliases should be snake_case if they are a value and CamelCase if they are a type. With the exception of the `DType` mapping types ex: `f32`. Alias names should clearly indicate what they are for and in addition to their docstring require no further information to understand assuming the reader understands the Mojo, and the domain. 18 | 19 | ## Functions 20 | 21 | Functions should be snake_case, and describe what they do in as few words as possible, such that in addition to the docstring no further info is required. 22 | 23 | The first line of a function docstring should summarize what the function does. 24 | ```python 25 | """ 26 | Description of the function. 27 | """ 28 | ``` 29 | Next add the parameters, arguments, and returns if there are any separated from the summary by a new line. For functions and parameters start with either `Parameters:` or `Args:` followed by a new line-separated list of the parameters or arguments with the name of the parameter/arg followed by a `:` and a description the description should be a sentence starting with a capital letter and ending with a period. For returns separated from previous lines by a new line and start with `Returns:` then go to a new line and write a brief description of the return value, again as a sentence starting with a capitol letter and ending with a period. If the function does not return the `Returns:` section should be omitted. 30 | 31 | There is no need to add the type name to the arguments or parameters as the compiler handles that. 32 | ```rust 33 | fn func[param:Copyable](arg1:param)->param: 34 | """ 35 | 36 | Description of the function. 37 | 38 | Parameters: 39 | param: Each parameter should be listed and described. 40 | 41 | Args: 42 | arg1: Each argument should be listed and described. 43 | 44 | Returns: 45 | Describe what is returned. 46 | """ 47 | ... 48 | ``` 49 | 50 | If the function has compile time constraints or raises `Error`s include sections similar to return that specify those constraints and possible errors. 51 | ```python 52 | """ 53 | Raises: 54 | A description of the errors raised by the function. 55 | 56 | Constraints: 57 | If the functions use compile time constraints they should be listed here. 58 | """ 59 | ``` 60 | 61 | ## Structs 62 | Structs should be CamelCase and describe what they do in as few words as possible, such that in addition to the docstring no further info is required. 63 | 64 | The first line of a struct docstring should summarize what the struct does. It is not necessary to reiterate the structs name in the docstring. The parameters, and constraints of a struct should be included in the struct docstring in a similar way to functions. 65 | 66 | ```rust 67 | struct AStruct[param:AnyType](AnyType): 68 | """ 69 | Struct docstring describes basically what a struct does. 70 | 71 | Constraints: 72 | Limitations placed on the struct. 73 | 74 | Parameters: 75 | param: An example parameter. 76 | """ 77 | ... 78 | ``` 79 | 80 | Fields and aliases should have a docstring below them describing what they are. They should be no longer than a single sentence and should start with a capital letter and end with a period. 81 | 82 | ```rust 83 | struct AStruct[param:AnyType](AnyType): 84 | """ 85 | Struct docstring describes basically what a struct does. 86 | 87 | Constraints: 88 | Limitations placed on the struct. 89 | 90 | Parameters: 91 | param: An example parameter. 92 | """ 93 | 94 | var field: Int64 95 | """ Field Descriptions go below each field.""" 96 | ... 97 | ``` 98 | 99 | Struct methods should follow the same rules as functions. 100 | 101 | ## Traits 102 | Traits follow the same rules as Structs but there are no fields in traits. 103 | -------------------------------------------------------------------------------- /docs/style.json: -------------------------------------------------------------------------------- 1 | { 2 | "decl": { 3 | "aliases": [ 4 | { 5 | "deprecated": "", 6 | "description": "", 7 | "kind": "alias", 8 | "name": "Example", 9 | "summary": "Aliases can be explained with docstrings and should if they exist in the global scope.", 10 | "value": "Int" 11 | } 12 | ], 13 | "description": "", 14 | "functions": [ 15 | { 16 | "kind": "function", 17 | "name": "func", 18 | "overloads": [ 19 | { 20 | "args": [ 21 | { 22 | "convention": "borrowed", 23 | "description": "Each arguament should be listed and described.", 24 | "kind": "argument", 25 | "name": "arg1", 26 | "passingKind": "pos_or_kw", 27 | "type": "param" 28 | } 29 | ], 30 | "async": false, 31 | "constraints": "If the functions use compile time constraints they should be listed here.", 32 | "deprecated": "", 33 | "description": "", 34 | "isDef": false, 35 | "isStatic": false, 36 | "kind": "function", 37 | "name": "func", 38 | "parameters": [ 39 | { 40 | "description": "Each parameter should be listed and described.", 41 | "kind": "parameter", 42 | "name": "param", 43 | "passingKind": "pos_or_kw", 44 | "type": "Copyable" 45 | } 46 | ], 47 | "raises": false, 48 | "raisesDoc": "", 49 | "returnType": "$0", 50 | "returnsDoc": "Describe what is returned.", 51 | "signature": "func[param: Copyable](arg1: param) -> $0", 52 | "summary": "Description of the function." 53 | } 54 | ] 55 | }, 56 | { 57 | "kind": "function", 58 | "name": "func1", 59 | "overloads": [ 60 | { 61 | "args": [ 62 | { 63 | "convention": "borrowed", 64 | "description": "Each arguament should be listed and described.", 65 | "kind": "argument", 66 | "name": "arg1", 67 | "passingKind": "pos_or_kw", 68 | "type": "param" 69 | } 70 | ], 71 | "async": false, 72 | "constraints": "", 73 | "deprecated": "", 74 | "description": "", 75 | "isDef": false, 76 | "isStatic": false, 77 | "kind": "function", 78 | "name": "func1", 79 | "parameters": [ 80 | { 81 | "description": "Each parameter should be listed and described.", 82 | "kind": "parameter", 83 | "name": "param", 84 | "passingKind": "pos_or_kw", 85 | "type": "Copyable" 86 | } 87 | ], 88 | "raises": true, 89 | "raisesDoc": "A description of the errors raised by the function.", 90 | "returnType": "$0", 91 | "returnsDoc": "Describe what is returned.", 92 | "signature": "func1[param: Copyable](arg1: param) -> $0", 93 | "summary": "Description of the function." 94 | } 95 | ] 96 | } 97 | ], 98 | "kind": "module", 99 | "name": "style", 100 | "structs": [ 101 | { 102 | "aliases": [], 103 | "constraints": "Limitations placed on the struct.", 104 | "deprecated": "", 105 | "description": "", 106 | "fields": [ 107 | { 108 | "description": "", 109 | "kind": "field", 110 | "name": "field", 111 | "summary": "Field Descriptions go below each field.", 112 | "type": "SIMD[int64, 1]" 113 | } 114 | ], 115 | "functions": [ 116 | { 117 | "kind": "function", 118 | "name": "func", 119 | "overloads": [ 120 | { 121 | "args": [ 122 | { 123 | "convention": "borrowed", 124 | "description": "", 125 | "kind": "argument", 126 | "name": "self", 127 | "passingKind": "pos_or_kw", 128 | "type": "Self" 129 | } 130 | ], 131 | "async": false, 132 | "constraints": "", 133 | "deprecated": "", 134 | "description": "", 135 | "isDef": false, 136 | "isStatic": false, 137 | "kind": "function", 138 | "name": "func", 139 | "parameters": [], 140 | "raises": false, 141 | "raisesDoc": "", 142 | "returnType": null, 143 | "returnsDoc": "", 144 | "signature": "func(self: Self)", 145 | "summary": "Function docstring like previosly shown." 146 | } 147 | ] 148 | } 149 | ], 150 | "kind": "struct", 151 | "name": "AStruct", 152 | "parameters": [ 153 | { 154 | "description": "An example parameter.", 155 | "kind": "parameter", 156 | "name": "param", 157 | "passingKind": "pos_or_kw", 158 | "type": "AnyType" 159 | } 160 | ], 161 | "parentTraits": [ 162 | "AnyType" 163 | ], 164 | "summary": "Struct docstring describing basically what a struct does." 165 | } 166 | ], 167 | "summary": "Document docstring decribing what it does, if it is in an init file it will be the docstring for the module", 168 | "traits": [ 169 | { 170 | "deprecated": "", 171 | "description": "", 172 | "fields": [], 173 | "functions": [ 174 | { 175 | "kind": "function", 176 | "name": "func", 177 | "overloads": [ 178 | { 179 | "args": [ 180 | { 181 | "convention": "borrowed", 182 | "description": "", 183 | "kind": "argument", 184 | "name": "self", 185 | "passingKind": "pos_or_kw", 186 | "type": "T" 187 | } 188 | ], 189 | "async": false, 190 | "constraints": "", 191 | "deprecated": "", 192 | "description": "", 193 | "isDef": false, 194 | "isStatic": false, 195 | "kind": "function", 196 | "name": "func", 197 | "parameters": [], 198 | "raises": false, 199 | "raisesDoc": "", 200 | "returnType": null, 201 | "returnsDoc": "", 202 | "signature": "func(self: T)", 203 | "summary": "Function docstring like previosly shown." 204 | } 205 | ] 206 | } 207 | ], 208 | "kind": "trait", 209 | "name": "ATrait", 210 | "parentTraits": [ 211 | "AnyType" 212 | ], 213 | "summary": "Describe the trait." 214 | } 215 | ] 216 | }, 217 | "version": "24.6.0" 218 | } -------------------------------------------------------------------------------- /docs/style_example.mojo: -------------------------------------------------------------------------------- 1 | """ 2 | Document docstring decribing what it does, if it is in an init file it will be the docstring for the module 3 | """ 4 | 5 | # ===----------------------------------------------------------------------=== # 6 | # Subsection header, used to divide code in a file by functional parts (ingored by doc generation) 7 | # ===----------------------------------------------------------------------=== # 8 | 9 | alias Example = Int 10 | """Aliases can be explained with docstrings and should if they exist in the global scope.""" 11 | 12 | 13 | fn func[param: Copyable](arg1: param) -> param: 14 | """ 15 | Description of the function. 16 | 17 | Constraints: 18 | If the functions use compile time constraints they should be listed here. 19 | 20 | Parameters: 21 | param: Each parameter should be listed and described. 22 | 23 | Args: 24 | arg1: Each arguament should be listed and described. 25 | 26 | Returns: 27 | Describe what is returned. 28 | """ 29 | return arg1 30 | 31 | 32 | fn func1[param: Copyable](arg1: param) raises -> param: 33 | """ 34 | Description of the function. 35 | 36 | Parameters: 37 | param: Each parameter should be listed and described. 38 | 39 | Args: 40 | arg1: Each arguament should be listed and described. 41 | 42 | Raises: 43 | A description of the errors raised by the function. 44 | 45 | Returns: 46 | Describe what is returned. 47 | """ 48 | return arg1 49 | 50 | 51 | struct AStruct[param: AnyType](AnyType): 52 | """ 53 | Struct docstring describing basically what a struct does. 54 | 55 | Constraints: 56 | Limitations placed on the struct. 57 | 58 | Parameters: 59 | param: An example parameter. 60 | """ 61 | 62 | var field: Int64 63 | """Field Descriptions go below each field.""" 64 | 65 | fn func(self) -> None: 66 | """ 67 | Function docstring like previosly shown. 68 | """ 69 | return None 70 | 71 | 72 | trait ATrait: 73 | """ 74 | Describe the trait. 75 | """ 76 | 77 | fn func(self) -> None: 78 | """ 79 | Function docstring like previosly shown. 80 | """ 81 | pass 82 | -------------------------------------------------------------------------------- /mojoproject.toml: -------------------------------------------------------------------------------- 1 | [project] 2 | name = "NuMojo" 3 | version = "0.6" 4 | description = "NuMojo is a library for numerical computing written in Mojo 🔥" 5 | authors = [ 6 | "Shivasankar ", 7 | "MadAlex1997 <>", 8 | "Yuhao Zhu (朱宇浩) ", 9 | "mmenendezg <>", 10 | "sandstromviktor <>", 11 | "durnwalder <>" 12 | ] 13 | channels = ["https://conda.modular.com/max", "https://repo.prefix.dev/modular-community", "conda-forge"] 14 | platforms = ["osx-arm64", "linux-64"] 15 | license = "Apache-2.0" 16 | readme = "README.MD" 17 | 18 | [tasks] 19 | # compile the package and copy it to the tests folder 20 | package = "magic run mojo package numojo && cp numojo.mojopkg tests/" 21 | p = "clear && magic run package" 22 | 23 | # format the package 24 | format = "magic run mojo format ./" 25 | 26 | # test whether tests pass on the built package 27 | test = "magic run package && magic run mojo test tests -I ./tests/" 28 | t = "clear && magic run test" 29 | 30 | # run individual tests to avoid overheat 31 | test_core = "magic run mojo test tests/core -I ./ -I ./tests/" 32 | test_creation = "magic run mojo test tests/routines/test_creation.mojo -I ./ -I ./tests/" 33 | test_functional = "magic run mojo test tests/routines/test_functional.mojo -I ./ -I ./tests/" 34 | test_indexing = "magic run mojo test tests/routines/test_indexing.mojo -I ./ -I ./tests/" 35 | test_linalg = "magic run mojo test tests/routines/test_linalg.mojo -I ./ -I ./tests/" 36 | test_manipulation = "magic run mojo test tests/routines/test_manipulation.mojo -I ./ -I ./tests/" 37 | test_math = "magic run mojo test tests/routines/test_math.mojo -I ./ -I ./tests/" 38 | test_random = "magic run mojo test tests/routines/test_random.mojo -I ./ -I ./tests/" 39 | test_statistics = "magic run mojo test tests/routines/test_statistics.mojo -I ./ -I ./tests/" 40 | test_sorting = "magic run mojo test tests/routines/test_sorting.mojo -I ./ -I ./tests/" 41 | 42 | # run all final checks before a commit 43 | final = "magic run format && magic run test" 44 | f = "clear && magic run final" 45 | 46 | # Automatically Generate doc pages 47 | doc_pages = "mojo doc numojo/ -o docs.json" 48 | 49 | # run everything and generate docs before release 50 | release = "clear && magic run final && magic run doc_pages" 51 | 52 | [dependencies] 53 | max = "=25.1.1" 54 | python = ">=3.11" 55 | numpy = ">=1.19" 56 | scipy = ">=1.14" -------------------------------------------------------------------------------- /numojo/__init__.mojo: -------------------------------------------------------------------------------- 1 | """ 2 | NuMojo is a library for numerical computing in Mojo 🔥 3 | similar to NumPy, SciPy in Python. 4 | """ 5 | 6 | alias __version__ = "V0.5" 7 | 8 | # ===----------------------------------------------------------------------=== # 9 | # Import core types 10 | # ===----------------------------------------------------------------------=== # 11 | 12 | from numojo.core.ndarray import NDArray 13 | from numojo.core.ndshape import NDArrayShape, Shape 14 | from numojo.core.ndstrides import NDArrayStrides, Strides 15 | from numojo.core.item import Item, item 16 | from numojo.core.complex.complex_dtype import CDType 17 | from numojo.core.complex.complex_simd import ComplexSIMD, ComplexScalar 18 | from numojo.core.complex.complex_ndarray import ComplexNDArray 19 | from numojo.core.matrix import Matrix 20 | from numojo.core.datatypes import ( 21 | i8, 22 | i16, 23 | i32, 24 | i64, 25 | isize, 26 | u8, 27 | u16, 28 | u32, 29 | u64, 30 | f16, 31 | f32, 32 | f64, 33 | ci8, 34 | ci16, 35 | ci32, 36 | ci64, 37 | cu8, 38 | cu16, 39 | cu32, 40 | cu64, 41 | cf16, 42 | cf32, 43 | cf64, 44 | ) 45 | 46 | # ===----------------------------------------------------------------------=== # 47 | # Import routines and objects 48 | # ===----------------------------------------------------------------------=== # 49 | 50 | # Objects 51 | from numojo.routines.constants import Constants 52 | 53 | alias pi = numojo.routines.constants.Constants.pi 54 | alias e = numojo.routines.constants.Constants.e 55 | alias c = numojo.routines.constants.Constants.c 56 | 57 | # Functions 58 | # TODO Make explicit imports of each individual function in future 59 | # to avoid polluting the root namespace. 60 | from numojo.routines import io 61 | from numojo.routines.io import ( 62 | loadtxt, 63 | savetxt, 64 | ) 65 | from numojo.routines.io import set_printoptions 66 | 67 | from numojo.routines import linalg 68 | from numojo.routines.linalg.misc import diagonal 69 | 70 | from numojo.routines import logic 71 | from numojo.routines.logic import ( 72 | greater, 73 | greater_equal, 74 | less, 75 | less_equal, 76 | equal, 77 | not_equal, 78 | isinf, 79 | isfinite, 80 | isnan, 81 | any, 82 | all, 83 | ) 84 | 85 | from numojo.routines import math 86 | from numojo.routines.math import ( 87 | add, 88 | sub, 89 | diff, 90 | mod, 91 | mul, 92 | div, 93 | floor_div, 94 | fma, 95 | remainder, 96 | ) 97 | from numojo.routines.math import gradient, trapz 98 | from numojo.routines.math import exp, exp2, expm1, log, ln, log2, log10, log1p 99 | from numojo.routines.math import ( 100 | max, 101 | min, 102 | mimimum, 103 | maximum, 104 | ) 105 | from numojo.routines.math import copysign 106 | from numojo.routines.math import ( 107 | arccosh, 108 | acosh, 109 | arcsinh, 110 | asinh, 111 | arctanh, 112 | atanh, 113 | cosh, 114 | sinh, 115 | tanh, 116 | ) 117 | from numojo.routines.math import cbrt, clip, rsqrt, sqrt, scalb 118 | from numojo.routines.math import prod, cumprod 119 | from numojo.routines.math import ( 120 | tabs, 121 | tfloor, 122 | tceil, 123 | ttrunc, 124 | tround, 125 | roundeven, 126 | nextafter, 127 | ) 128 | from numojo.routines.math import sum, cumsum 129 | from numojo.routines.math import ( 130 | arccos, 131 | acos, 132 | arcsin, 133 | asin, 134 | arctan, 135 | atan, 136 | atan2, 137 | cos, 138 | sin, 139 | tan, 140 | hypot, 141 | hypot_fma, 142 | ) 143 | 144 | from numojo.routines import statistics 145 | from numojo.routines.statistics import ( 146 | mean, 147 | mode, 148 | median, 149 | variance, 150 | std, 151 | ) 152 | 153 | from numojo.routines import bitwise 154 | from numojo.routines.bitwise import invert 155 | 156 | from numojo.routines import creation 157 | from numojo.routines.creation import ( 158 | arange, 159 | linspace, 160 | logspace, 161 | geomspace, 162 | empty, 163 | empty_like, 164 | eye, 165 | identity, 166 | ones, 167 | ones_like, 168 | zeros, 169 | zeros_like, 170 | full, 171 | full_like, 172 | diag, 173 | diagflat, 174 | tri, 175 | tril, 176 | triu, 177 | vander, 178 | fromstring, 179 | from_tensor, 180 | array, 181 | ) 182 | 183 | from numojo.routines import indexing 184 | from numojo.routines.indexing import where, compress, take_along_axis 185 | 186 | from numojo.routines.functional import apply_along_axis 187 | 188 | from numojo.routines import manipulation 189 | from numojo.routines.manipulation import ( 190 | ndim, 191 | shape, 192 | size, 193 | reshape, 194 | ravel, 195 | transpose, 196 | broadcast_to, 197 | flip, 198 | ) 199 | 200 | from numojo.routines import random 201 | 202 | from numojo.routines import sorting 203 | from numojo.routines.sorting import sort, argsort 204 | 205 | from numojo.routines import searching 206 | from numojo.routines.searching import argmax, argmin 207 | 208 | # ===----------------------------------------------------------------------=== # 209 | # Alias for users 210 | # For ease of use, the name of the types may not follow the Mojo convention, 211 | # e.g., lower case can also be used for alias of structs. 212 | # ===----------------------------------------------------------------------=== # 213 | -------------------------------------------------------------------------------- /numojo/core/__init__.mojo: -------------------------------------------------------------------------------- 1 | # ARRAYS 2 | 3 | # Structs 4 | from .ndarray import NDArray 5 | from .item import Item 6 | from .ndshape import NDArrayShape 7 | from .ndstrides import NDArrayStrides 8 | 9 | from .complex import ( 10 | CDType, 11 | ComplexSIMD, 12 | ComplexScalar, 13 | ComplexNDArray, 14 | ) 15 | 16 | from .datatypes import ( 17 | i8, 18 | i16, 19 | i32, 20 | i64, 21 | u8, 22 | u16, 23 | u32, 24 | u64, 25 | f16, 26 | f32, 27 | f64, 28 | ci8, 29 | ci16, 30 | ci32, 31 | ci64, 32 | cu8, 33 | cu16, 34 | cu32, 35 | cu64, 36 | cf16, 37 | cf32, 38 | cf64, 39 | ) 40 | 41 | # from .utility import 42 | 43 | alias idx = Item 44 | alias shape = NDArrayShape 45 | alias Shape = NDArrayShape 46 | -------------------------------------------------------------------------------- /numojo/core/_array_funcs.mojo: -------------------------------------------------------------------------------- 1 | # from ..traits.NDArrayTraits import NDArrayBackend 2 | from algorithm.functional import parallelize, vectorize, num_physical_cores 3 | from sys import simdwidthof 4 | 5 | from numojo.core.ndarray import NDArray 6 | 7 | """ 8 | Implementing backend for array keeping it simple for now 9 | """ 10 | 11 | 12 | fn math_func_1_array_in_one_array_out[ 13 | dtype: DType, 14 | func: fn[type: DType, simd_w: Int] (SIMD[type, simd_w]) -> SIMD[ 15 | type, simd_w 16 | ], 17 | ](array: NDArray[dtype]) raises -> NDArray[dtype]: 18 | """ 19 | Apply a SIMD compatible function to a NDArray and returns a new NDArray. 20 | 21 | Parameters: 22 | dtype: The NDArray element type. 23 | func: The SIMD compatible function to act on the NDArray. 24 | 25 | Args: 26 | array: A NDArray. 27 | 28 | Returns: 29 | A new NDArray that is the result of applying the function to the NDArray. 30 | """ 31 | var result_array: NDArray[dtype] = NDArray[dtype](array.shape) 32 | alias width = simdwidthof[dtype]() 33 | 34 | @parameter 35 | fn closure[simd_width: Int](i: Int): 36 | var simd_data = array._buf.ptr.load[width=simd_width](i) 37 | result_array._buf.ptr.store(i, func[dtype, simd_width](simd_data)) 38 | 39 | vectorize[closure, width](array.size) 40 | 41 | return result_array 42 | 43 | 44 | fn math_func_2_array_in_one_array_out[ 45 | dtype: DType, 46 | func: fn[type: DType, simd_w: Int] ( 47 | SIMD[type, simd_w], SIMD[type, simd_w] 48 | ) -> SIMD[type, simd_w], 49 | ](array1: NDArray[dtype], array2: NDArray[dtype]) raises -> NDArray[dtype]: 50 | """ 51 | Apply a SIMD compatible function to two NDArrays and returns a new NDArray. 52 | 53 | Raises: 54 | Error if the two arrays do not have the same shape. 55 | 56 | Parameters: 57 | dtype: The NDArray element type. 58 | func: The SIMD compatible function to act on the NDArrays. 59 | 60 | Args: 61 | array1: A NDArray. 62 | array2: A NDArray. 63 | 64 | Returns: 65 | A new NDArray that is the result of applying the function to the input NDArrays. 66 | """ 67 | 68 | if array1.shape != array2.shape: 69 | raise Error("Shape Mismatch error shapes must match for this function") 70 | 71 | var result_array: NDArray[dtype] = NDArray[dtype](array1.shape) 72 | alias width = simdwidthof[dtype]() 73 | 74 | @parameter 75 | fn closure[simd_width: Int](i: Int): 76 | var simd_data1 = array1._buf.ptr.load[width=simd_width](i) 77 | var simd_data2 = array2._buf.ptr.load[width=simd_width](i) 78 | result_array._buf.ptr.store( 79 | i, func[dtype, simd_width](simd_data1, simd_data2) 80 | ) 81 | 82 | vectorize[closure, width](result_array.size) 83 | 84 | return result_array 85 | 86 | 87 | fn math_func_one_array_one_SIMD_in_one_array_out[ 88 | dtype: DType, 89 | func: fn[type: DType, simd_w: Int] ( 90 | SIMD[type, simd_w], SIMD[type, simd_w] 91 | ) -> SIMD[type, simd_w], 92 | ](array: NDArray[dtype], scalar: SIMD[dtype, 1]) raises -> NDArray[dtype]: 93 | """ 94 | Apply a SIMD compatible function to a NDArray and a SIMD value and returns a new NDArray. 95 | 96 | Parameters: 97 | dtype: The NDArray element type. 98 | func: The SIMD compatible function to act on the NDArray and SIMD value. 99 | 100 | Args: 101 | array: A NDArray. 102 | scalar: A scalar value. 103 | 104 | Returns: 105 | A new NDArray that is the result of applying the function to the input NDArray and SIMD value. 106 | """ 107 | 108 | var result_array: NDArray[dtype] = NDArray[dtype](array.shape) 109 | alias width = simdwidthof[dtype]() 110 | 111 | @parameter 112 | fn closure[simd_width: Int](i: Int): 113 | var simd_data1 = array._buf.ptr.load[width=simd_width](i) 114 | result_array._buf.ptr.store( 115 | i, func[dtype, simd_width](simd_data1, scalar) 116 | ) 117 | 118 | vectorize[closure, width](result_array.size) 119 | return result_array 120 | -------------------------------------------------------------------------------- /numojo/core/complex/__init__.mojo: -------------------------------------------------------------------------------- 1 | from .complex_dtype import CDType 2 | from .complex_simd import ComplexSIMD, ComplexScalar 3 | from .complex_ndarray import ComplexNDArray 4 | -------------------------------------------------------------------------------- /numojo/core/datatypes.mojo: -------------------------------------------------------------------------------- 1 | """ 2 | Datatypes Module - Implements datatypes aliases, conversions 3 | """ 4 | # ===----------------------------------------------------------------------=== # 5 | # Datatypes Module - Implements datatypes aliases, conversions 6 | # Last updated: 2024-06-18 7 | # ===----------------------------------------------------------------------=== # 8 | 9 | # Rust-like or numpy-like data type alias 10 | """alias for `DType.int8`""" 11 | alias i8 = DType.int8 12 | """Data type alias for DType.int8""" 13 | alias i16 = DType.int16 14 | """Data type alias for DType.int16""" 15 | alias i32 = DType.int32 16 | """Data type alias for DType.int32""" 17 | alias i64 = DType.int64 18 | """Data type alias for DType.int64""" 19 | alias isize = DType.index 20 | """Data type alias for DType.index""" 21 | alias intp = DType.index 22 | """Data type alias for DType.index""" 23 | alias u8 = DType.uint8 24 | """Data type alias for DType.uint8""" 25 | alias u16 = DType.uint16 26 | """Data type alias for DType.uint16""" 27 | alias u32 = DType.uint32 28 | """Data type alias for DType.uint32""" 29 | alias u64 = DType.uint64 30 | """Data type alias for DType.uint64""" 31 | alias f16 = DType.float16 32 | """Data type alias for DType.float16""" 33 | alias f32 = DType.float32 34 | """Data type alias for DType.float32""" 35 | alias f64 = DType.float64 36 | """Data type alias for DType.float64""" 37 | alias boolean = DType.bool 38 | """Data type alias for DType.bool""" 39 | 40 | # ===----------------------------------------------------------------------=== # 41 | 42 | # Complex SIMD data type aliases 43 | """ Data type alias for ComplexSIMD[DType.int8, 1] """ 44 | alias ci8 = CDType.int8 45 | """ Data type alias for ComplexSIMD[DType.int16, 1] """ 46 | alias ci16 = CDType.int16 47 | """ Data type alias for ComplexSIMD[DType.int32, 1] """ 48 | alias ci32 = CDType.int32 49 | """ Data type alias for ComplexSIMD[DType.int64, 1] """ 50 | alias ci64 = CDType.int64 51 | """ Data type alias for ComplexSIMD[DType.uint8, 1] """ 52 | alias cu8 = CDType.uint8 53 | """ Data type alias for ComplexSIMD[DType.uint16, 1] """ 54 | alias cu16 = CDType.uint16 55 | """ Data type alias for ComplexSIMD[DType.uint32, 1] """ 56 | alias cu32 = CDType.uint32 57 | """ Data type alias for ComplexSIMD[DType.uint64, 1] """ 58 | alias cu64 = CDType.uint64 59 | """ Data type alias for ComplexSIMD[DType.float16, 1] """ 60 | alias cf16 = CDType.float16 61 | """ Data type alias for ComplexSIMD[DType.float32, 1] """ 62 | alias cf32 = CDType.float32 63 | """ Data type alias for ComplexSIMD[DType.float64, 1] """ 64 | alias cf64 = CDType.float64 65 | 66 | # ===----------------------------------------------------------------------=== # 67 | 68 | 69 | # TODO: Optimize the conditions with dict and move it to compile time 70 | # Dict can't be created at compile time rn 71 | struct TypeCoercion: 72 | """Handles type coercion using a promotion matrix approach.""" 73 | 74 | alias ranks: List[DType] = List[DType]( 75 | i8, u8, f16, i16, u16, f32, i32, u32, i64, u64, f64 76 | ) 77 | alias int_ranks: List[DType] = List[DType]( 78 | i8, u8, i16, u16, i32, u32, i64, u64 79 | ) 80 | alias float_ranks: List[DType] = List[DType](f16, f32, f64) 81 | 82 | @parameter 83 | @staticmethod 84 | fn get_type_rank[dtype: DType]() -> Int: 85 | try: 86 | return Self.ranks.index(dtype) 87 | except ValueError: 88 | return 10 89 | 90 | @parameter 91 | @staticmethod 92 | fn get_inttype_rank[dtype: DType]() -> Int: 93 | try: 94 | return Self.int_ranks.index(dtype) 95 | except ValueError: 96 | return 7 97 | 98 | @parameter 99 | @staticmethod 100 | fn get_floattype_rank[dtype: DType]() -> Int: 101 | try: 102 | return Self.float_ranks.index(dtype) 103 | except ValueError: 104 | return 2 105 | 106 | @parameter 107 | @staticmethod 108 | fn coerce_floats[T1: DType, T2: DType]() -> DType: 109 | """Coerces two floating point types.""" 110 | if T1 == f16 or T2 == f16: 111 | if T1 == f64 or T2 == f64: 112 | return f64 113 | return f32 114 | var rank1 = Self.get_floattype_rank[T1]() 115 | var rank2 = Self.get_floattype_rank[T2]() 116 | return T1 if rank1 > rank2 else T2 117 | 118 | @parameter 119 | @staticmethod 120 | fn coerce_signed_ints[T1: DType, T2: DType]() -> DType: 121 | """Coerces two signed integer types.""" 122 | var rank1 = Self.get_type_rank[T1]() 123 | var rank2 = Self.get_type_rank[T2]() 124 | var max_rank = max(rank1, rank2) 125 | if max_rank <= 3: 126 | return i16 # int8 -> int16 127 | if max_rank <= 6: 128 | return i32 # int16 -> int32 129 | if max_rank <= 8: 130 | return i64 # int32 -> int64 131 | return f64 # int64 -> float64 132 | 133 | @parameter 134 | @staticmethod 135 | fn coerce_unsigned_ints[T1: DType, T2: DType]() -> DType: 136 | """Coerces two unsigned integer types.""" 137 | if T1.sizeof() >= T2.sizeof(): 138 | return T1 139 | else: 140 | return T2 141 | 142 | @parameter 143 | @staticmethod 144 | fn coerce_mixed_ints[T1: DType, T2: DType]() -> DType: 145 | """Coerces a signed and unsigned integer type.""" 146 | alias signed = T1 if T1.is_signed() else T2 147 | alias unsigned = T2 if T1.is_signed() else T1 148 | 149 | # Handle signed/unsigned pairs 150 | if signed == i8 and unsigned == u8: 151 | return i16 152 | if signed == i16 and unsigned == u16: 153 | return i32 154 | if signed == i32 and unsigned == u32: 155 | return i64 156 | if signed == i64 and unsigned == u64: 157 | return f64 158 | 159 | # If unsigned type is larger, use next larger signed type 160 | var signed_rank = Self.get_type_rank[signed]() 161 | var unsigned_rank = Self.get_type_rank[unsigned]() 162 | 163 | if unsigned_rank > signed_rank: 164 | if unsigned == u16: 165 | return i32 166 | if unsigned == u32: 167 | return i64 168 | if unsigned == u64: 169 | return f64 170 | 171 | return signed 172 | 173 | @parameter 174 | @staticmethod 175 | fn coerce_mixed[int_type: DType, float_type: DType]() -> DType: 176 | """Coerces a mixed integer and floating point type.""" 177 | # Special case: float16 always promotes to at least float32 178 | if float_type == f16 and (int_type == i16 or int_type == u16): 179 | return f32 180 | if float_type == f16 and (int_type == i32 or int_type == u32): 181 | return f64 182 | if float_type == f16 and (int_type == i64 or int_type == u64): 183 | return f64 184 | # Special cases for int32/uint32 and larger with float32 185 | if float_type == f32: 186 | if int_type in (i32, u32, i64, u64): 187 | return f64 188 | return float_type 189 | 190 | @parameter 191 | @staticmethod 192 | fn result[T1: DType, T2: DType]() -> DType: 193 | """Returns the coerced output type for two input types.""" 194 | if T1 == T2: 195 | return T1 196 | elif T1.is_floating_point() and T2.is_floating_point(): 197 | return TypeCoercion.coerce_floats[T1, T2]() 198 | elif T1.is_integral() and T2.is_integral(): 199 | if T1.is_signed() and T2.is_signed(): 200 | return TypeCoercion.coerce_signed_ints[T1, T2]() 201 | elif T1.is_unsigned() and T2.is_unsigned(): 202 | return TypeCoercion.coerce_unsigned_ints[T1, T2]() 203 | else: 204 | return TypeCoercion.coerce_mixed_ints[T1, T2]() 205 | elif T1.is_integral() and T2.is_floating_point(): 206 | return TypeCoercion.coerce_mixed[T1, T2]() 207 | elif T1.is_floating_point() and T2.is_integral(): 208 | return TypeCoercion.coerce_mixed[T2, T1]() 209 | return T1 210 | 211 | 212 | fn _concise_dtype_str(dtype: DType) -> String: 213 | """Returns a concise string representation of the data type.""" 214 | if dtype == i8: 215 | return "i8" 216 | elif dtype == i16: 217 | return "i16" 218 | elif dtype == i32: 219 | return "i32" 220 | elif dtype == i64: 221 | return "i64" 222 | elif dtype == isize: 223 | return "index" 224 | elif dtype == u8: 225 | return "u8" 226 | elif dtype == u16: 227 | return "u16" 228 | elif dtype == u32: 229 | return "u32" 230 | elif dtype == u64: 231 | return "u64" 232 | elif dtype == f16: 233 | return "f16" 234 | elif dtype == f32: 235 | return "f32" 236 | elif dtype == f64: 237 | return "f64" 238 | elif dtype == boolean: 239 | return "boolean" 240 | elif dtype == isize: 241 | return "isize" 242 | else: 243 | return "Unknown" 244 | -------------------------------------------------------------------------------- /numojo/core/flags.mojo: -------------------------------------------------------------------------------- 1 | # ===----------------------------------------------------------------------=== # 2 | # Distributed under the Apache 2.0 License with LLVM Exceptions. 3 | # See LICENSE and the LLVM License for more information. 4 | # https://github.com/Mojo-Numerics-and-Algorithms-group/NuMojo/blob/main/LICENSE 5 | # https://llvm.org/LICENSE.txt 6 | # ===----------------------------------------------------------------------=== # 7 | """ 8 | Implements Flags type. 9 | """ 10 | 11 | from numojo.core.ndshape import NDArrayShape 12 | from numojo.core.ndstrides import NDArrayStrides 13 | 14 | 15 | @register_passable 16 | struct Flags: 17 | """ 18 | Information about the memory layout of the array. 19 | The Flags object can be accessed dictionary-like. 20 | or by using lowercased attribute names. 21 | Short names are available for convenience when using dictionary-like access. 22 | """ 23 | 24 | # attributes 25 | var C_CONTIGUOUS: Bool 26 | """C_CONTIGUOUS (C): The data is in a C-style contiguous segment.""" 27 | var F_CONTIGUOUS: Bool 28 | """F_CONTIGUOUS (F): The data is in a Fortran-style contiguous segment.""" 29 | var OWNDATA: Bool 30 | """OWNDATA (O): The array owns the underlying data buffer.""" 31 | var WRITEABLE: Bool 32 | """ 33 | The data area can be written to. 34 | If it is False, the data is read-only and be blocked from writing. 35 | The WRITEABLE field of a view or slice is inherited from the array where 36 | it is derived. If the parent object is not writeable, the child object is 37 | also not writeable. If the parent object is writeable, the child object may 38 | be not writeable. 39 | """ 40 | var FORC: Bool 41 | """F_CONTIGUOUS or C_CONTIGUOUS.""" 42 | 43 | # === ---------------------------------------------------------------- === # 44 | # Life cycle dunder methods 45 | # === ---------------------------------------------------------------- === # 46 | 47 | fn __init__( 48 | out self, 49 | c_contiguous: Bool, 50 | f_contiguous: Bool, 51 | owndata: Bool, 52 | writeable: Bool, 53 | ): 54 | """ 55 | Initializes the Flags object with provided information. 56 | 57 | Args: 58 | c_contiguous: The data is in a C-style contiguous segment. 59 | f_contiguous: The data is in a Fortran-style contiguous segment. 60 | owndata: The array owns the underlying data buffer. 61 | writeable: The data area can be written to. 62 | If owndata is False, writeable is forced to be False. 63 | """ 64 | 65 | self.C_CONTIGUOUS = c_contiguous 66 | self.F_CONTIGUOUS = f_contiguous 67 | self.OWNDATA = owndata 68 | self.WRITEABLE = writeable and owndata 69 | self.FORC = f_contiguous or c_contiguous 70 | 71 | fn __init__( 72 | out self, 73 | shape: NDArrayShape, 74 | strides: NDArrayStrides, 75 | owndata: Bool, 76 | writeable: Bool, 77 | ) raises: 78 | """ 79 | Initializes the Flags object according the shape and strides information. 80 | 81 | Args: 82 | shape: The shape of the array. 83 | strides: The strides of the array. 84 | owndata: The array owns the underlying data buffer. 85 | writeable: The data area can be written to. 86 | If owndata is False, writeable is forced to be False. 87 | """ 88 | 89 | self.C_CONTIGUOUS = ( 90 | True if (strides[-1] == 1) or (shape[-1] == 1) else False 91 | ) 92 | self.F_CONTIGUOUS = ( 93 | True if (strides[0] == 1) or (shape[0] == 1) else False 94 | ) 95 | self.OWNDATA = owndata 96 | self.WRITEABLE = writeable and owndata 97 | self.FORC = self.F_CONTIGUOUS or self.C_CONTIGUOUS 98 | 99 | fn __init__( 100 | out self, 101 | shape: Tuple[Int, Int], 102 | strides: Tuple[Int, Int], 103 | owndata: Bool, 104 | writeable: Bool, 105 | ): 106 | """ 107 | Initializes the Flags object according the shape and strides information. 108 | 109 | Args: 110 | shape: The shape of the array. 111 | strides: The strides of the array. 112 | owndata: The array owns the underlying data buffer. 113 | writeable: The data area can be written to. 114 | If owndata is False, writeable is forced to be False. 115 | """ 116 | 117 | self.C_CONTIGUOUS = ( 118 | True if (strides[1] == 1) or (shape[1] == 1) else False 119 | ) 120 | self.F_CONTIGUOUS = ( 121 | True if (strides[0] == 1) or (shape[0] == 1) else False 122 | ) 123 | self.OWNDATA = owndata 124 | self.WRITEABLE = writeable and owndata 125 | self.FORC = self.F_CONTIGUOUS or self.C_CONTIGUOUS 126 | 127 | fn __copyinit__(out self, other: Self): 128 | """ 129 | Initializes the Flags object by copying the information from 130 | another Flags object. 131 | 132 | Args: 133 | other: The Flags object to copy information from. 134 | """ 135 | 136 | self.C_CONTIGUOUS = other.C_CONTIGUOUS 137 | self.F_CONTIGUOUS = other.F_CONTIGUOUS 138 | self.OWNDATA = other.OWNDATA 139 | self.WRITEABLE = other.WRITEABLE 140 | self.FORC = other.FORC 141 | 142 | # === ---------------------------------------------------------------- === # 143 | # Get and set dunder methods 144 | # === ---------------------------------------------------------------- === # 145 | 146 | fn __getitem__(self, key: String) raises -> Bool: 147 | """ 148 | Get the value of the fields with the given key. 149 | The Flags object can be accessed dictionary-like. 150 | Short names are available for convenience. 151 | 152 | Args: 153 | key: The key of the field to get. 154 | 155 | Returns: 156 | The value of the field with the given key. 157 | """ 158 | if ( 159 | (key != "C_CONTIGUOUS") 160 | and (key != "C") 161 | and (key != "F_CONTIGUOUS") 162 | and (key != "F") 163 | and (key != "OWNDATA") 164 | and (key != "O") 165 | and (key != "WRITEABLE") 166 | and (key != "W") 167 | and (key != "FORC") 168 | ): 169 | raise Error( 170 | String( 171 | "\nError in `Flags.__getitem__()`: " 172 | "Invalid field name or short name: {}".format(key) 173 | ) 174 | ) 175 | if (key == "C_CONTIGUOUS") or (key == "C"): 176 | return self.C_CONTIGUOUS 177 | elif (key == "F_CONTIGUOUS") or (key == "F"): 178 | return self.F_CONTIGUOUS 179 | elif (key == "OWNDATA") or (key == "O"): 180 | return self.OWNDATA 181 | elif (key == "WRITEABLE") or (key == "W"): 182 | return self.WRITEABLE 183 | else: 184 | return self.FORC 185 | -------------------------------------------------------------------------------- /numojo/core/own_data.mojo: -------------------------------------------------------------------------------- 1 | # ===----------------------------------------------------------------------=== # 2 | # Define `OwnData` type 3 | # 4 | # TODO: fields in traits are not supported yet by Mojo 5 | # Currently use `get_ptr()` to get pointer, in future, use `ptr` directly. 6 | # var ptr: UnsafePointer[Scalar[dtype]] 7 | # TODO: implement `Bufferable` trait. 8 | # ===----------------------------------------------------------------------=== # 9 | 10 | from memory import UnsafePointer 11 | from numojo.core.traits.bufferable import Bufferable 12 | 13 | 14 | struct OwnData[dtype: DType]: # TODO: implement `Bufferable` trait 15 | var ptr: UnsafePointer[Scalar[dtype]] 16 | 17 | fn __init__(out self, size: Int): 18 | """ 19 | Allocate given space on memory. 20 | The bytes allocated is `size` * `byte size of dtype`. 21 | 22 | Notes: 23 | `ndarray.flags['OWN_DATA']` should be set as True. 24 | The memory should be freed by `__del__`. 25 | """ 26 | self.ptr = UnsafePointer[Scalar[dtype]]().alloc(size) 27 | 28 | fn __init__(out self, ptr: UnsafePointer[Scalar[dtype]]): 29 | """ 30 | Do not use this if you know what it means. 31 | If the pointer is associated with another array, it might cause 32 | dangling pointer problem. 33 | 34 | Notes: 35 | `ndarray.flags['OWN_DATA']` should be set as False. 36 | The memory should not be freed by `__del__`. 37 | """ 38 | self.ptr = ptr 39 | 40 | fn __moveinit__(out self, owned other: Self): 41 | self.ptr = other.ptr 42 | 43 | fn get_ptr(self) -> UnsafePointer[Scalar[dtype]]: 44 | return self.ptr 45 | -------------------------------------------------------------------------------- /numojo/core/ref_data.mojo: -------------------------------------------------------------------------------- 1 | # ===----------------------------------------------------------------------=== # 2 | # Define `RefData` type 3 | # 4 | # TODO: fields in traits are not supported yet by Mojo 5 | # Currently use `get_ptr()` to get pointer, in future, use `ptr` directly. 6 | # var ptr: UnsafePointer[Float16] 7 | # TODO: use parameterized trait. 8 | # Replace `Float16` with `Scalar[dtype]` 9 | # ===----------------------------------------------------------------------=== # 10 | 11 | from memory import UnsafePointer 12 | from numojo.core.traits.bufferable import Bufferable 13 | 14 | 15 | struct RefData[is_mutable: Bool, //, origin: Origin[is_mutable]](Bufferable): 16 | var ptr: UnsafePointer[Float16] 17 | 18 | fn __init__(out self, size: Int): 19 | """ 20 | Allocate given space on memory. 21 | The bytes allocated is `size` * `byte size of dtype`. 22 | 23 | Notes: 24 | Although it has the lifetime of another array, it owns the data. 25 | `ndarray.flags['OWN_DATA']` should be set as True. 26 | The memory should be freed by `__del__`. 27 | """ 28 | self.ptr = UnsafePointer[Float16]().alloc(size) 29 | 30 | fn __init__(out self, ptr: UnsafePointer[Float16]): 31 | """ 32 | Reads the underlying data of another array. 33 | 34 | Notes: 35 | `ndarray.flags['OWN_DATA']` should be set as False. 36 | The memory should not be freed by `__del__`. 37 | """ 38 | self.ptr = ptr 39 | 40 | fn __moveinit__(out self, owned other: Self): 41 | self.ptr = other.ptr 42 | 43 | fn get_ptr(self) -> UnsafePointer[Float16]: 44 | return self.ptr 45 | -------------------------------------------------------------------------------- /numojo/core/traits/__init__.mojo: -------------------------------------------------------------------------------- 1 | """ 2 | Defines Numojo Traits 3 | """ 4 | from .backend import Backend 5 | -------------------------------------------------------------------------------- /numojo/core/traits/array_like.mojo: -------------------------------------------------------------------------------- 1 | from numojo.core.ndarray import NDArray 2 | 3 | # Blocked by lack of trait paramaterization 4 | 5 | # trait Arraylike: 6 | # fn load[width: Int](self, idx: Int) -> SIMD[dtype, width]: 7 | # """ 8 | # Loads a SIMD element of size `width` at the given index `idx`. 9 | # """ 10 | # ... 11 | # fn store[width: Int](mut self, idx: Int, val: SIMD[dtype, width]): 12 | # """ 13 | # Stores the SIMD element of size `width` at index `idx`. 14 | # """ 15 | # ... 16 | 17 | # trait NDArrayBackend: 18 | # """ 19 | # A trait that defines backends for calculations in the rest of the library. 20 | # """ 21 | 22 | # fn __init__(mut self: Self): 23 | # """ 24 | # Initialize the backend. 25 | # """ 26 | # ... 27 | 28 | # fn math_func_1_array_in_one_array_out[ 29 | # dtype: DType, 30 | # func: fn[type: DType, simd_w: Int] (SIMD[type, simd_w]) -> SIMD[ 31 | # type, simd_w 32 | # ], 33 | # ](self: Self, array: Arraylike) -> Arraylike: 34 | # """ 35 | # Apply a SIMD function of one variable and one return to a NDArray 36 | 37 | # Parameters: 38 | # dtype: The element type. 39 | # func: the SIMD function to to apply. 40 | 41 | # Args: 42 | # array: A NDArray 43 | 44 | # Returns: 45 | # A a new NDArray that is NDArray with the function func applied. 46 | # """ 47 | # ... 48 | 49 | # fn math_func_2_array_in_one_array_out[ 50 | # dtype: DType, 51 | # func: fn[type: DType, simd_w: Int] ( 52 | # SIMD[type, simd_w], SIMD[type, simd_w] 53 | # ) -> SIMD[type, simd_w], 54 | # ]( 55 | # self: Self, array1: Arraylike, array2: Arraylike 56 | # ) raises -> Arraylike: 57 | # """ 58 | # Apply a SIMD function of two variable and one return to a NDArray 59 | 60 | # Constraints: 61 | # Both arrays must have the same shape 62 | 63 | # Parameters: 64 | # dtype: The element type. 65 | # func: the SIMD function to to apply. 66 | 67 | # Args: 68 | # array1: A NDArray 69 | # array2: A NDArray 70 | 71 | # Returns: 72 | # A a new NDArray that is NDArray with the function func applied. 73 | # """ 74 | 75 | # ... 76 | 77 | # fn math_func_one_array_one_SIMD_in_one_array_out[ 78 | # dtype: DType, 79 | # func: fn[type: DType, simd_w: Int] ( 80 | # SIMD[type, simd_w], SIMD[type, simd_w] 81 | # ) -> SIMD[type, simd_w], 82 | # ]( 83 | # self: Self, array: Arraylike, scalar: Scalar[dtype] 84 | # ) -> Arraylike: 85 | # """ 86 | # Apply a SIMD function of two variable and one return to a NDArray 87 | 88 | # Constraints: 89 | # Both arrays must have the same shape 90 | 91 | # Parameters: 92 | # dtype: The element type. 93 | # func: the SIMD function to to apply. 94 | 95 | # Args: 96 | # array: A NDArray 97 | # scalar: A Scalar 98 | 99 | # Returns: 100 | # A a new NDArray that is NDArray with the function func applied. 101 | # """ 102 | 103 | # ... 104 | -------------------------------------------------------------------------------- /numojo/core/traits/backend.mojo: -------------------------------------------------------------------------------- 1 | # ===----------------------------------------------------------------------=== # 2 | # Defines computational backend traits 3 | # ===----------------------------------------------------------------------=== # 4 | 5 | 6 | from numojo.core.ndarray import NDArray 7 | 8 | 9 | trait Backend: 10 | """ 11 | A trait that defines backends for calculations in the rest of the library. 12 | """ 13 | 14 | fn __init__(mut self: Self): 15 | """ 16 | Initialize the backend. 17 | """ 18 | pass 19 | 20 | fn math_func_fma[ 21 | dtype: DType, 22 | ]( 23 | self: Self, 24 | array1: NDArray[dtype], 25 | array2: NDArray[dtype], 26 | array3: NDArray[dtype], 27 | ) raises -> NDArray[dtype]: 28 | """ 29 | Apply a SIMD level fuse multipy add function of three variables and one return to a NDArray. 30 | 31 | Constraints: 32 | Both arrays must have the same shape 33 | 34 | Parameters: 35 | dtype: The element type. 36 | 37 | Args: 38 | array1: A NDArray. 39 | array2: A NDArray. 40 | array3: A NDArray. 41 | 42 | Returns: 43 | A a new NDArray that is NDArray with the function func applied. 44 | 45 | Raises: 46 | If shapes are missmatched or there is a access error. 47 | """ 48 | pass 49 | 50 | fn math_func_fma[ 51 | dtype: DType, 52 | ]( 53 | self: Self, 54 | array1: NDArray[dtype], 55 | array2: NDArray[dtype], 56 | simd: SIMD[dtype, 1], 57 | ) raises -> NDArray[dtype]: 58 | """ 59 | Apply a SIMD level fuse multipy add function of three variables and one return to a NDArray. 60 | 61 | Constraints: 62 | Both arrays must have the same shape 63 | 64 | Parameters: 65 | dtype: The element type. 66 | 67 | Args: 68 | array1: A NDArray. 69 | array2: A NDArray. 70 | simd: A SIMD[dtype,1] value to be added. 71 | 72 | Returns: 73 | A new NDArray that is NDArray with the function func applied. 74 | """ 75 | pass 76 | 77 | fn math_func_1_array_in_one_array_out[ 78 | dtype: DType, 79 | func: fn[type: DType, simd_w: Int] (SIMD[type, simd_w]) -> SIMD[ 80 | type, simd_w 81 | ], 82 | ](self: Self, array: NDArray[dtype]) raises -> NDArray[dtype]: 83 | """ 84 | Apply a SIMD function of one variable and one return to a NDArray. 85 | 86 | Parameters: 87 | dtype: The element type. 88 | func: The SIMD function to to apply. 89 | 90 | Args: 91 | array: A NDArray. 92 | 93 | Returns: 94 | A new NDArray that is NDArray with the function func applied. 95 | """ 96 | ... 97 | 98 | fn math_func_2_array_in_one_array_out[ 99 | dtype: DType, 100 | func: fn[type: DType, simd_w: Int] ( 101 | SIMD[type, simd_w], SIMD[type, simd_w] 102 | ) -> SIMD[type, simd_w], 103 | ]( 104 | self: Self, array1: NDArray[dtype], array2: NDArray[dtype] 105 | ) raises -> NDArray[dtype]: 106 | """ 107 | Apply a SIMD function of two variable and one return to a NDArray. 108 | 109 | Constraints: 110 | Both arrays must have the same shape 111 | 112 | Parameters: 113 | dtype: The element type. 114 | func: The SIMD function to to apply. 115 | 116 | Args: 117 | array1: A NDArray. 118 | array2: A NDArray. 119 | 120 | Returns: 121 | A new NDArray that is NDArray with the function func applied. 122 | """ 123 | 124 | ... 125 | 126 | fn math_func_1_array_1_scalar_in_one_array_out[ 127 | dtype: DType, 128 | func: fn[type: DType, simd_w: Int] ( 129 | SIMD[type, simd_w], SIMD[type, simd_w] 130 | ) -> SIMD[type, simd_w], 131 | ]( 132 | self: Self, array: NDArray[dtype], scalar: Scalar[dtype] 133 | ) raises -> NDArray[dtype]: 134 | """ 135 | Apply a SIMD function of two variable and one return to a NDArray. 136 | 137 | Constraints: 138 | Both arrays must have the same shape 139 | 140 | Parameters: 141 | dtype: The element type. 142 | func: The SIMD function to to apply. 143 | 144 | Args: 145 | array: A NDArray. 146 | scalar: A Scalars. 147 | 148 | Returns: 149 | A new NDArray that is NDArray with the function func applied. 150 | """ 151 | 152 | ... 153 | 154 | fn math_func_compare_2_arrays[ 155 | dtype: DType, 156 | func: fn[type: DType, simd_w: Int] ( 157 | SIMD[type, simd_w], SIMD[type, simd_w] 158 | ) -> SIMD[DType.bool, simd_w], 159 | ]( 160 | self: Self, array1: NDArray[dtype], array2: NDArray[dtype] 161 | ) raises -> NDArray[DType.bool]: 162 | """ 163 | Apply a SIMD comparision function of two variable. 164 | 165 | Constraints: 166 | Both arrays must have the same shape. 167 | 168 | Parameters: 169 | dtype: The element type. 170 | func: The SIMD comparision function to to apply. 171 | 172 | Args: 173 | array1: A NDArray. 174 | array2: A NDArray. 175 | 176 | Returns: 177 | A new Boolean NDArray that is NDArray with the function func applied. 178 | """ 179 | ... 180 | 181 | fn math_func_compare_array_and_scalar[ 182 | dtype: DType, 183 | func: fn[type: DType, simd_w: Int] ( 184 | SIMD[type, simd_w], SIMD[type, simd_w] 185 | ) -> SIMD[DType.bool, simd_w], 186 | ]( 187 | self: Self, array1: NDArray[dtype], scalar: SIMD[dtype, 1] 188 | ) raises -> NDArray[DType.bool]: 189 | """ 190 | Apply a SIMD comparision function of two variable. 191 | 192 | Constraints: 193 | Both arrays must have the same shape. 194 | 195 | Parameters: 196 | dtype: The element type. 197 | func: The SIMD comparision function to to apply. 198 | 199 | Args: 200 | array1: A NDArray. 201 | scalar: A scalar. 202 | 203 | Returns: 204 | A new Boolean NDArray that is NDArray with the function func applied. 205 | """ 206 | ... 207 | 208 | fn math_func_is[ 209 | dtype: DType, 210 | func: fn[type: DType, simd_w: Int] (SIMD[type, simd_w]) -> SIMD[ 211 | DType.bool, simd_w 212 | ], 213 | ](self: Self, array: NDArray[dtype]) raises -> NDArray[DType.bool]: 214 | ... 215 | 216 | # fn math_func_simd_int[ 217 | # dtype: DType, 218 | # func: fn[type: DType, simd_w: Int] (SIMD[type, simd_w], Int) -> SIMD[ 219 | # type, simd_w 220 | # ], 221 | # ](self: Self, array1: NDArray[dtype], intval: Int) raises -> NDArray[dtype]: 222 | # ... 223 | -------------------------------------------------------------------------------- /numojo/core/traits/bufferable.mojo: -------------------------------------------------------------------------------- 1 | # ===----------------------------------------------------------------------=== # 2 | # Define `Bufferable` traits 3 | # ===----------------------------------------------------------------------=== # 4 | 5 | from memory import UnsafePointer 6 | 7 | 8 | trait Bufferable: 9 | """ 10 | Data buffer types that can be used as a container of the underlying buffer. 11 | """ 12 | 13 | # TODO: fields in traits are not supported yet by Mojo 14 | # Currently use `get_ptr()` to get pointer, in future, use `ptr` directly. 15 | # var ptr: UnsafePointer[Float16] 16 | # TODO: use parameterized trait. 17 | # Replace `Float16` with `Scalar[dtype]` 18 | 19 | fn __init__(out self, size: Int): 20 | ... 21 | 22 | fn __init__(out self, ptr: UnsafePointer[Float16]): 23 | ... 24 | 25 | fn __moveinit__(out self, owned other: Self): 26 | ... 27 | 28 | fn get_ptr(self) -> UnsafePointer[Float16]: 29 | ... 30 | -------------------------------------------------------------------------------- /numojo/core/traits/indexer_collection_element.mojo: -------------------------------------------------------------------------------- 1 | trait IndexerCollectionElement(CollectionElement, Indexer): 2 | """The IndexerCollectionElement trait denotes a trait composition 3 | of the `Indexer` and `CollectionElement` traits. 4 | 5 | This is useful to have as a named entity since Mojo does not 6 | currently support anonymous trait compositions to constrain 7 | on `Indexer & CollectionElement` in the parameter. 8 | """ 9 | 10 | pass 11 | -------------------------------------------------------------------------------- /numojo/prelude.mojo: -------------------------------------------------------------------------------- 1 | """ 2 | prelude 3 | ======= 4 | 5 | NuMojo comes a wide range of functions, types, and constants. 6 | If you manually import everything, 7 | it will make the header of the file too long. 8 | On the other hand, using `from numojo import *` would import a lot of functions 9 | that you never use and would pollute the naming space. 10 | 11 | This module tries to find out a balance by providing a list of things 12 | that can be imported at one time. 13 | The list contains the functions or types 14 | that are the most essential for a user. 15 | 16 | You can use the following code to import them: 17 | 18 | ```mojo 19 | from numojo.prelude import * 20 | ``` 21 | """ 22 | 23 | import numojo as nm 24 | 25 | from numojo.core.item import Item, item 26 | from numojo.core.matrix import Matrix 27 | from numojo.core.ndarray import NDArray 28 | from numojo.core.ndshape import Shape, NDArrayShape 29 | 30 | from numojo.core.complex.complex_dtype import CDType 31 | from numojo.core.complex.complex_simd import ComplexSIMD, ComplexScalar 32 | from numojo.core.complex.complex_ndarray import ComplexNDArray 33 | 34 | from numojo.core.datatypes import ( 35 | i8, 36 | i16, 37 | i32, 38 | i64, 39 | isize, 40 | intp, 41 | u8, 42 | u16, 43 | u32, 44 | u64, 45 | f16, 46 | f32, 47 | f64, 48 | boolean, 49 | ci8, 50 | ci16, 51 | ci32, 52 | ci64, 53 | cu8, 54 | cu16, 55 | cu32, 56 | cu64, 57 | cf16, 58 | cf32, 59 | cf64, 60 | ) 61 | -------------------------------------------------------------------------------- /numojo/routines/__init__.mojo: -------------------------------------------------------------------------------- 1 | """ 2 | Implements routines by topic: 3 | 4 | - Array creation routines (creation.mojo) 5 | - Array manipulation routines (manipulation.mojo) 6 | - Bit-wise operations (bitwise.mojo) 7 | - Constants (constants.mojo) 8 | - Indexing routines (indexing.mojo) 9 | - Input and output (io/) 10 | - Text files (files.mojo) 11 | - Text formatting options (formatting.mojo) 12 | - Linear algebra (linalg/) 13 | - Decompositions (decompositions.mojo) 14 | - Products of matrices and vectors (products.mojo) 15 | - Solving (solving.mojo) 16 | - Logic functions (logic/) 17 | - Comparison (comparison.mojo) 18 | - Array contents (contents.mojo) 19 | - Truth value testing (truth.mojo) 20 | - Mathematical functions (math/) 21 | - Arithmetic operations (arithmetic.mojo) 22 | - Exponents and logarithms (exponents.mojo) 23 | - Extrema finding (extrema.mojo) 24 | - Floating point routines (floating.mojo) 25 | - Hyperbolic functions (hyper.mojo) 26 | - Miscellaneous (misc.mojo) 27 | - Rounding (rounding.mojo) 28 | - Sums, products, differences (sums.mojo, products.mojo, differences.mojo) 29 | - Trigonometric functions (trig.mojo) 30 | - Random sampling (random.mojo) 31 | - Sorting, searching, and counting (sorting.mojo, searching.mojo) 32 | - Statistics (statistics/) 33 | - Averages and variances (averages.mojo) 34 | 35 | """ 36 | -------------------------------------------------------------------------------- /numojo/routines/bitwise.mojo: -------------------------------------------------------------------------------- 1 | # ===------------------------------------------------------------------------===# 2 | # Bit-wise operations 3 | # ===------------------------------------------------------------------------===# 4 | 5 | # ===------------------------------------------------------------------------===# 6 | # Element-wise bit operations 7 | # ===------------------------------------------------------------------------===# 8 | 9 | 10 | import math 11 | from algorithm import parallelize 12 | from algorithm import Static2DTileUnitFunc as Tile2DFunc 13 | from utils import Variant 14 | 15 | import numojo.core._math_funcs as _mf 16 | from numojo.core.ndarray import NDArray, NDArrayShape 17 | from numojo.core.utility import is_inttype, is_floattype, is_booltype 18 | 19 | 20 | fn invert[ 21 | dtype: DType, backend: _mf.Backend = _mf.Vectorized 22 | ](array: NDArray[dtype]) raises -> NDArray[dtype]: 23 | """ 24 | Element-wise invert of an array. 25 | 26 | Constraints: 27 | The array must be either a boolean or integral array. 28 | 29 | Parameters: 30 | dtype: The element type. 31 | backend: Sets utility function origin, defaults to `Vectorized`. 32 | 33 | Args: 34 | array: A NDArray. 35 | 36 | Returns: 37 | A NDArray equal to the bitwise inversion of array. 38 | """ 39 | constrained[ 40 | is_inttype[dtype]() or is_booltype[dtype](), 41 | "Only Bools and integral types can be invertedd.", 42 | ]() 43 | 44 | return backend().math_func_1_array_in_one_array_out[dtype, SIMD.__invert__]( 45 | array 46 | ) 47 | -------------------------------------------------------------------------------- /numojo/routines/constants.mojo: -------------------------------------------------------------------------------- 1 | """ 2 | Constants 3 | """ 4 | # ===----------------------------------------------------------------------=== # 5 | # Implements Constants 6 | # Last updated: 2024-06-16 7 | # ===----------------------------------------------------------------------=== # 8 | 9 | 10 | @value 11 | struct Constants(AnyType): 12 | """Define constants. 13 | 14 | Use alias for compile time evaluation of indefinite precision. 15 | ```mojo 16 | import numojo as nm 17 | fn main(): 18 | var pi: Float64 = nm.pi 19 | print("Float64:", pi*pi*pi*pi*pi*pi) 20 | print("Literal:", nm.pi*nm.pi*nm.pi*nm.pi*nm.pi*nm.pi) 21 | ``` 22 | ```console 23 | Float64: 961.38919357530415 24 | Literal: 961.38919357530449 25 | ``` 26 | """ 27 | 28 | alias c = 299_792_458 29 | alias pi = 3.1415926535897932384626433832795028841971693937510582097494459230781640628620899862803482534211706798214808651328230664709384460955058223172535940812848111745028410270193852110555954930381966446229489 30 | alias e = 2.71828182845904523536028747135266249775724609375 31 | alias hbar = 1.0545718176461563912626e-34 32 | 33 | fn __init__(mut self): 34 | """ 35 | Initializes the constants. 36 | """ 37 | pass 38 | 39 | fn __del__(owned self): 40 | """ 41 | Deletes the constants. 42 | """ 43 | pass 44 | -------------------------------------------------------------------------------- /numojo/routines/io/__init__.mojo: -------------------------------------------------------------------------------- 1 | from .files import ( 2 | loadtxt, 3 | savetxt, 4 | ) 5 | 6 | from .formatting import ( 7 | format_floating_scientific, 8 | PrintOptions, 9 | set_printoptions, 10 | ) 11 | -------------------------------------------------------------------------------- /numojo/routines/io/files.mojo: -------------------------------------------------------------------------------- 1 | from numojo.routines.creation import fromstring 2 | from collections.optional import Optional 3 | 4 | 5 | # contains a custom basic implementation of loadtxt and savetxt to be used temporarily 6 | # until the official implementation is ready 7 | # one could use numpy backend, but it might add a dependency to numpy 8 | # better load files through numpy and then pass it to Numojo through array() function 9 | fn loadtxt[ 10 | dtype: DType = f64 11 | ]( 12 | filename: String, 13 | delimiter: String = ",", 14 | skiprows: Int = 0, 15 | usecols: Optional[List[Int]] = None, 16 | ) raises -> NDArray[dtype]: 17 | with open(filename, "r") as file: 18 | string = file.read() 19 | var shape_offset_init: Int = string.find("[") 20 | var shape_offset_fin: Int = string.find("]") 21 | var ndim_offset_init: Int = string.find("[", start=shape_offset_fin) 22 | var ndim_offset_fin: Int = string.find("]", start=ndim_offset_init) 23 | var ndim: Int = Int(string[ndim_offset_init + 1 : ndim_offset_fin]) 24 | var ndshape: List[Int] = List[Int]() 25 | for i in range(shape_offset_init + 1, shape_offset_fin): 26 | if string[i].isdigit(): 27 | ndshape.append(Int(string[i])) 28 | var data: List[Scalar[dtype]] = List[Scalar[dtype]]() 29 | for i in range(ndim_offset_fin + 2, len(string)): 30 | if string[i].isdigit(): 31 | var number: String = string[i] 32 | data.append(atof(number).cast[dtype]()) 33 | return array[dtype](data=data, shape=ndshape, order="C") 34 | 35 | 36 | fn savetxt[ 37 | dtype: DType = f64 38 | ](filename: String, array: NDArray[dtype], delimiter: String = ",") raises: 39 | var shape: String = "ndshape=[" 40 | for i in range(array.ndim): 41 | shape += String(array.shape[i]) 42 | if i != array.ndim - 1: 43 | shape = shape + ", " 44 | shape = shape + "]" 45 | print(shape) 46 | 47 | with open(filename, "w") as file: 48 | file.write(shape + "\n") 49 | file.write("ndim=[" + String(array.ndim) + "]\n") 50 | for i in range(array.size): 51 | if i % 10 == 0: 52 | file.write(String("\n")) 53 | file.write(String(array._buf.ptr[i]) + ",") 54 | -------------------------------------------------------------------------------- /numojo/routines/linalg/__init__.mojo: -------------------------------------------------------------------------------- 1 | from .decompositions import lu_decomposition, qr 2 | from .norms import det, trace 3 | from .products import cross, dot, matmul 4 | from .solving import inv, solve, lstsq 5 | from .misc import diagonal 6 | -------------------------------------------------------------------------------- /numojo/routines/linalg/misc.mojo: -------------------------------------------------------------------------------- 1 | # ===----------------------------------------------------------------------=== # 2 | # Distributed under the Apache 2.0 License with LLVM Exceptions. 3 | # See LICENSE and the LLVM License for more information. 4 | # https://github.com/Mojo-Numerics-and-Algorithms-group/NuMojo/blob/main/LICENSE 5 | # https://llvm.org/LICENSE.txt 6 | # ===----------------------------------------------------------------------=== # 7 | 8 | # ===----------------------------------------------------------------------=== # 9 | # Miscellaneous Linear Algebra Routines 10 | # ===----------------------------------------------------------------------=== # 11 | 12 | from numojo.core.ndarray import NDArray 13 | 14 | 15 | fn diagonal[ 16 | dtype: DType 17 | ](a: NDArray[dtype], offset: Int = 0) raises -> NDArray[dtype]: 18 | """ 19 | Returns specific diagonals. 20 | Currently supports only 2D arrays. 21 | 22 | Raises: 23 | Error: If the array is not 2D. 24 | Error: If the offset is beyond the shape of the array. 25 | 26 | Parameters: 27 | dtype: Data type of the array. 28 | 29 | Args: 30 | a: An NDArray. 31 | offset: Offset of the diagonal from the main diagonal. 32 | 33 | Returns: 34 | The diagonal of the NDArray. 35 | """ 36 | 37 | if a.ndim != 2: 38 | raise Error("\nError in `diagonal`: Only supports 2D arrays") 39 | 40 | var m = a.shape[0] 41 | var n = a.shape[1] 42 | 43 | if offset >= max(m, n): # Offset beyond the shape of the array 44 | raise Error( 45 | "\nError in `diagonal`: Offset beyond the shape of the array" 46 | ) 47 | 48 | var res: NDArray[dtype] 49 | 50 | if offset >= 0: 51 | var size_of_res = min(n - offset, m) 52 | res = NDArray[dtype](Shape(size_of_res)) 53 | for i in range(size_of_res): 54 | res.item(i) = a.item(i, i + offset) 55 | else: 56 | var size_of_res = min(m + offset, m) 57 | res = NDArray[dtype](Shape(size_of_res)) 58 | for i in range(size_of_res): 59 | res.item(i) = a.item(i - offset, i) 60 | 61 | return res 62 | -------------------------------------------------------------------------------- /numojo/routines/linalg/norms.mojo: -------------------------------------------------------------------------------- 1 | # ===----------------------------------------------------------------------=== # 2 | # Norms and other numbers 3 | # ===----------------------------------------------------------------------=== # 4 | 5 | from numojo.core.ndarray import NDArray 6 | from numojo.core.matrix import Matrix 7 | from numojo.routines.linalg.decompositions import ( 8 | lu_decomposition, 9 | partial_pivoting, 10 | ) 11 | 12 | 13 | fn det[dtype: DType](A: NDArray[dtype]) raises -> Scalar[dtype]: 14 | """ 15 | Find the determinant of A using LUP decomposition. 16 | """ 17 | 18 | if A.ndim != 2: 19 | raise Error(String("Array must be 2d.")) 20 | if A.shape[0] != A.shape[1]: 21 | raise Error(String("Matrix is not square.")) 22 | 23 | var det_L: Scalar[dtype] = 1 24 | var det_U: Scalar[dtype] = 1 25 | var n = A.shape[0] # Dimension of the matrix 26 | 27 | var A_pivoted: NDArray[dtype] 28 | var U: NDArray[dtype] 29 | var L: NDArray[dtype] 30 | var s: Int 31 | A_pivoted, _, s = partial_pivoting(A) 32 | L, U = lu_decomposition[dtype](A_pivoted) 33 | 34 | for i in range(n): 35 | det_L = det_L * L.item(i, i) 36 | det_U = det_U * U.item(i, i) 37 | 38 | if s % 2 == 0: 39 | return det_L * det_U 40 | else: 41 | return -det_L * det_U 42 | 43 | 44 | fn det[dtype: DType](A: Matrix[dtype]) raises -> Scalar[dtype]: 45 | """ 46 | Find the determinant of A using LUP decomposition. 47 | """ 48 | var det_L: Scalar[dtype] = 1 49 | var det_U: Scalar[dtype] = 1 50 | var n = A.shape[0] # Dimension of the matrix 51 | 52 | var U: Matrix[dtype] 53 | var L: Matrix[dtype] 54 | A_pivoted, _, s = partial_pivoting(A) 55 | L, U = lu_decomposition[dtype](A_pivoted) 56 | 57 | for i in range(n): 58 | det_L = det_L * L[i, i] 59 | det_U = det_U * U[i, i] 60 | 61 | if s % 2 == 0: 62 | return det_L * det_U 63 | else: 64 | return -det_L * det_U 65 | 66 | 67 | # TODO: implement for arbitrary axis 68 | fn trace[ 69 | dtype: DType 70 | ]( 71 | array: NDArray[dtype], offset: Int = 0, axis1: Int = 0, axis2: Int = 1 72 | ) raises -> NDArray[dtype]: 73 | """ 74 | Computes the trace of a ndarray. 75 | 76 | Parameters: 77 | dtype: Data type of the array. 78 | 79 | Args: 80 | array: A NDArray. 81 | offset: Offset of the diagonal from the main diagonal. 82 | axis1: First axis. 83 | axis2: Second axis. 84 | 85 | Returns: 86 | The trace of the NDArray. 87 | """ 88 | if array.ndim != 2: 89 | raise Error("Trace is currently only supported for 2D arrays") 90 | if axis1 > array.ndim - 1 or axis2 > array.ndim - 1: 91 | raise Error("axis cannot be greater than the rank of the array") 92 | var result: NDArray[dtype] = NDArray[dtype](Shape(1)) 93 | var rows = array.shape[0] 94 | var cols = array.shape[1] 95 | var diag_length = min(rows, cols - offset) if offset >= 0 else min( 96 | rows + offset, cols 97 | ) 98 | 99 | for i in range(diag_length): 100 | var row = i if offset >= 0 else i - offset 101 | var col = i + offset if offset >= 0 else i 102 | result._buf.ptr.store( 103 | 0, result._buf.ptr.load(0) + array._buf.ptr[row * cols + col] 104 | ) 105 | 106 | return result 107 | 108 | 109 | fn trace[ 110 | dtype: DType 111 | ](A: Matrix[dtype], offset: Int = 0) raises -> Scalar[dtype]: 112 | """ 113 | Return the sum along diagonals of the array. 114 | 115 | Similar to `numpy.trace`. 116 | """ 117 | var m = A.shape[0] 118 | var n = A.shape[1] 119 | 120 | if offset >= max(m, n): # Offset beyond the shape of the matrix 121 | return 0 122 | 123 | var res = Scalar[dtype](0) 124 | 125 | if offset >= 0: 126 | for i in range(n - offset): 127 | res = res + A[i, i + offset] 128 | else: 129 | for i in range(m + offset): 130 | res = res + A[i - offset, i] 131 | 132 | return res 133 | -------------------------------------------------------------------------------- /numojo/routines/logic/__init__.mojo: -------------------------------------------------------------------------------- 1 | from .comparison import ( 2 | greater, 3 | greater_equal, 4 | less, 5 | less_equal, 6 | equal, 7 | not_equal, 8 | ) 9 | from .contents import isinf, isfinite, isnan 10 | from .truth import any, all 11 | -------------------------------------------------------------------------------- /numojo/routines/logic/contents.mojo: -------------------------------------------------------------------------------- 1 | """ 2 | Implements Checking routines: currently not SIMD due to bool bit packing issue 3 | """ 4 | # ===----------------------------------------------------------------------=== # 5 | # Array contents 6 | # ===----------------------------------------------------------------------=== # 7 | 8 | 9 | import math 10 | 11 | import numojo.core._math_funcs as _mf 12 | from numojo.core.ndarray import NDArray 13 | 14 | # fn is_power_of_2[ 15 | # dtype: DType, backend: _mf.Backend = _mf.Vectorized 16 | # ](array: NDArray[dtype]) -> NDArray[DType.bool]: 17 | # return backend().math_func_is[dtype, math.is_power_of_2](array) 18 | 19 | 20 | # fn is_even[ 21 | # dtype: DType, backend: _mf.Backend = _mf.Vectorized 22 | # ](array: NDArray[dtype]) -> NDArray[DType.bool]: 23 | # return backend().math_func_is[dtype, math.is_even](array) 24 | 25 | 26 | # fn is_odd[ 27 | # dtype: DType, backend: _mf.Backend = _mf.Vectorized 28 | # ](array: NDArray[dtype]) -> NDArray[DType.bool]: 29 | # return backend().math_func_is[dtype, math.is_odd](array) 30 | 31 | 32 | fn isinf[ 33 | dtype: DType, backend: _mf.Backend = _mf.Vectorized 34 | ](array: NDArray[dtype]) raises -> NDArray[DType.bool]: 35 | """ 36 | Checks if each element of the input array is infinite. 37 | 38 | Parameters: 39 | dtype: DType - Data type of the input array. 40 | backend: _mf.Backend - Backend to use for the operation. Defaults to _mf.Vectorized. 41 | 42 | Args: 43 | array: NDArray[dtype] - Input array to check. 44 | 45 | Returns: 46 | NDArray[DType.bool] - A array of the same shape as `array` with True for infinite elements and False for others. 47 | """ 48 | # return backend().math_func_is[dtype, math.isinf](array) 49 | 50 | var result_array: NDArray[DType.bool] = NDArray[DType.bool](array.shape) 51 | for i in range(result_array.size): 52 | result_array.store(i, math.isinf(array.load(i))) 53 | return result_array 54 | 55 | 56 | fn isfinite[ 57 | dtype: DType, backend: _mf.Backend = _mf.Vectorized 58 | ](array: NDArray[dtype]) raises -> NDArray[DType.bool]: 59 | """ 60 | Checks if each element of the input array is finite. 61 | 62 | Parameters: 63 | dtype: DType - Data type of the input array. 64 | backend: _mf.Backend - Backend to use for the operation. Defaults to _mf.Vectorized. 65 | 66 | Args: 67 | array: NDArray[dtype] - Input array to check. 68 | 69 | Returns: 70 | NDArray[DType.bool] - A array of the same shape as `array` with True for finite elements and False for others. 71 | """ 72 | # return backend().math_func_is[dtype, math.isfinite](array) 73 | var result_array: NDArray[DType.bool] = NDArray[DType.bool](array.shape) 74 | for i in range(result_array.size): 75 | result_array.store(i, math.isfinite(array.load(i))) 76 | return result_array 77 | 78 | 79 | fn isnan[ 80 | dtype: DType, backend: _mf.Backend = _mf.Vectorized 81 | ](array: NDArray[dtype]) raises -> NDArray[DType.bool]: 82 | """ 83 | Checks if each element of the input array is NaN. 84 | 85 | Parameters: 86 | dtype: DType - Data type of the input array. 87 | backend: _mf.Backend - Backend to use for the operation. Defaults to _mf.Vectorized. 88 | 89 | Args: 90 | array: NDArray[dtype] - Input array to check. 91 | 92 | Returns: 93 | NDArray[DType.bool] - A array of the same shape as `array` with True for NaN elements and False for others. 94 | """ 95 | # return backend().math_func_is[dtype, math.isnan](array) 96 | var result_array: NDArray[DType.bool] = NDArray[DType.bool](array.shape) 97 | for i in range(result_array.size): 98 | result_array.store(i, math.isnan(array.load(i))) 99 | return result_array 100 | -------------------------------------------------------------------------------- /numojo/routines/logic/truth.mojo: -------------------------------------------------------------------------------- 1 | # ===----------------------------------------------------------------------=== # 2 | # Truth value testing 3 | # ===----------------------------------------------------------------------=== # 4 | 5 | import math 6 | from algorithm import vectorize, parallelize 7 | from sys import simdwidthof 8 | 9 | import numojo.core._math_funcs as _mf 10 | from numojo.core.ndarray import NDArray 11 | from numojo.core.matrix import Matrix 12 | 13 | 14 | fn all[dtype: DType](A: Matrix[dtype]) -> Scalar[dtype]: 15 | """ 16 | Test whether all array elements evaluate to True. 17 | 18 | Args: 19 | A: Matrix. 20 | """ 21 | var res = Scalar[dtype](1) 22 | alias width: Int = simdwidthof[dtype]() 23 | 24 | @parameter 25 | fn cal_and[width: Int](i: Int): 26 | res = res & A._buf.ptr.load[width=width](i).reduce_and() 27 | 28 | vectorize[cal_and, width](A.size) 29 | return res 30 | 31 | 32 | fn all[dtype: DType](A: Matrix[dtype], axis: Int) raises -> Matrix[dtype]: 33 | """ 34 | Test whether all array elements evaluate to True along axis. 35 | """ 36 | 37 | alias width: Int = simdwidthof[dtype]() 38 | 39 | if axis == 0: 40 | var B = Matrix.ones[dtype](shape=(1, A.shape[1])) 41 | 42 | for i in range(A.shape[0]): 43 | 44 | @parameter 45 | fn cal_vec_sum[width: Int](j: Int): 46 | B._store[width]( 47 | 0, j, B._load[width](0, j) & A._load[width](i, j) 48 | ) 49 | 50 | vectorize[cal_vec_sum, width](A.shape[1]) 51 | 52 | return B^ 53 | 54 | elif axis == 1: 55 | var B = Matrix.ones[dtype](shape=(A.shape[0], 1)) 56 | 57 | @parameter 58 | fn cal_rows(i: Int): 59 | @parameter 60 | fn cal_sum[width: Int](j: Int): 61 | B._store( 62 | i, 63 | 0, 64 | B._load(i, 0) & A._load[width=width](i, j).reduce_and(), 65 | ) 66 | 67 | vectorize[cal_sum, width](A.shape[1]) 68 | 69 | parallelize[cal_rows](A.shape[0], A.shape[0]) 70 | return B^ 71 | 72 | else: 73 | raise Error(String("The axis can either be 1 or 0!")) 74 | 75 | 76 | fn allt(array: NDArray[DType.bool]) raises -> Scalar[DType.bool]: 77 | """ 78 | If all True. 79 | 80 | Args: 81 | array: A NDArray. 82 | Returns: 83 | A boolean scalar 84 | """ 85 | var result = Scalar[DType.bool](True) 86 | # alias opt_nelts: Int = simdwidthof[DType.bool]() 87 | 88 | # @parameter 89 | # fn vectorize_sum[simd_width: Int](idx: Int) -> None: 90 | # var simd_data = array.load[width=simd_width](idx) 91 | # result |= simd_data.reduce_and() 92 | 93 | # vectorize[vectorize_sum, opt_nelts](array.size) 94 | # return result 95 | for i in range(array.size): 96 | result &= array.load(i) 97 | return result 98 | 99 | 100 | fn any(array: NDArray[DType.bool]) raises -> Scalar[DType.bool]: 101 | """ 102 | If any True. 103 | 104 | Args: 105 | array: A NDArray. 106 | Returns: 107 | A boolean scalar 108 | """ 109 | var result = Scalar[DType.bool](False) 110 | # alias opt_nelts: Int = simdwidthof[DType.bool]() 111 | 112 | # @parameter 113 | # fn vectorize_sum[simd_width: Int](idx: Int) -> None: 114 | # var simd_data = array.load[width=simd_width](idx) 115 | # result &= simd_data.reduce_or() 116 | 117 | # vectorize[vectorize_sum, opt_nelts](array.size) 118 | # return result 119 | for i in range(array.size): 120 | result |= array.load(i) 121 | return result 122 | 123 | 124 | fn any[dtype: DType](A: Matrix[dtype]) -> Scalar[dtype]: 125 | """ 126 | Test whether any array elements evaluate to True. 127 | 128 | Args: 129 | A: Matrix. 130 | """ 131 | var res = Scalar[dtype](0) 132 | alias width: Int = simdwidthof[dtype]() 133 | 134 | @parameter 135 | fn cal_and[width: Int](i: Int): 136 | res = res | A._buf.ptr.load[width=width](i).reduce_or() 137 | 138 | vectorize[cal_and, width](A.size) 139 | return res 140 | 141 | 142 | fn any[dtype: DType](A: Matrix[dtype], axis: Int) raises -> Matrix[dtype]: 143 | """ 144 | Test whether any array elements evaluate to True along axis. 145 | """ 146 | 147 | alias width: Int = simdwidthof[dtype]() 148 | 149 | if axis == 0: 150 | var B = Matrix.zeros[dtype](shape=(1, A.shape[1])) 151 | 152 | for i in range(A.shape[0]): 153 | 154 | @parameter 155 | fn cal_vec_sum[width: Int](j: Int): 156 | B._store[width]( 157 | 0, j, B._load[width](0, j) | A._load[width](i, j) 158 | ) 159 | 160 | vectorize[cal_vec_sum, width](A.shape[1]) 161 | 162 | return B^ 163 | 164 | elif axis == 1: 165 | var B = Matrix.zeros[dtype](shape=(A.shape[0], 1)) 166 | 167 | @parameter 168 | fn cal_rows(i: Int): 169 | @parameter 170 | fn cal_sum[width: Int](j: Int): 171 | B._store( 172 | i, 173 | 0, 174 | B._load(i, 0) | A._load[width=width](i, j).reduce_or(), 175 | ) 176 | 177 | vectorize[cal_sum, width](A.shape[1]) 178 | 179 | parallelize[cal_rows](A.shape[0], A.shape[0]) 180 | return B^ 181 | 182 | else: 183 | raise Error(String("The axis can either be 1 or 0!")) 184 | -------------------------------------------------------------------------------- /numojo/routines/math/__init__.mojo: -------------------------------------------------------------------------------- 1 | from .arithmetic import add, sub, diff, mod, mul, div, floor_div, fma, remainder 2 | from .differences import gradient, trapz 3 | from .exponents import exp, exp2, expm1, log, ln, log2, log10, log1p 4 | from .extrema import max, min, mimimum, maximum 5 | from .floating import copysign 6 | from .hyper import ( 7 | arccosh, 8 | acosh, 9 | arcsinh, 10 | asinh, 11 | arctanh, 12 | atanh, 13 | cosh, 14 | sinh, 15 | tanh, 16 | ) 17 | from .misc import cbrt, clip, rsqrt, sqrt, scalb 18 | from .products import prod, cumprod 19 | from .rounding import ( 20 | round, 21 | tabs, 22 | tfloor, 23 | tceil, 24 | ttrunc, 25 | tround, 26 | roundeven, 27 | nextafter, 28 | ) 29 | from .sums import sum, cumsum 30 | from .trig import ( 31 | arccos, 32 | acos, 33 | arcsin, 34 | asin, 35 | arctan, 36 | atan, 37 | atan2, 38 | cos, 39 | sin, 40 | tan, 41 | hypot, 42 | hypot_fma, 43 | ) 44 | -------------------------------------------------------------------------------- /numojo/routines/math/differences.mojo: -------------------------------------------------------------------------------- 1 | # ===----------------------------------------------------------------------=== # 2 | # Differences 3 | # ===----------------------------------------------------------------------=== # 4 | 5 | import math 6 | from algorithm import parallelize 7 | from algorithm import Static2DTileUnitFunc as Tile2DFunc 8 | 9 | import numojo.core._math_funcs as _mf 10 | from numojo.routines.creation import arange 11 | from numojo.core.ndarray import NDArray 12 | from numojo.core.utility import is_inttype, is_floattype 13 | 14 | """TODO: 15 | 1) add a Variant[NDArray, Scalar, ...] to include all possibilities 16 | 2) add edge_order 17 | """ 18 | 19 | 20 | fn gradient[ 21 | dtype: DType = DType.float64 22 | ](x: NDArray[dtype], spacing: Scalar[dtype]) raises -> NDArray[dtype]: 23 | """ 24 | Compute the gradient of y over x using the trapezoidal rule. 25 | 26 | Parameters: 27 | dtype: Input data type. 28 | 29 | Args: 30 | x: An array. 31 | spacing: An array of the same shape as x containing the spacing between adjacent elements. 32 | 33 | Constraints: 34 | `fdtype` must be a floating-point type if `idtype` is not a floating-point type. 35 | 36 | Returns: 37 | The integral of y over x using the trapezoidal rule. 38 | """ 39 | 40 | var result: NDArray[dtype] = NDArray[dtype](x.shape) 41 | var space: NDArray[dtype] = arange[dtype](1, x.size + 1, step=spacing) 42 | var hu: Scalar[dtype] = space.load(1) 43 | var hd: Scalar[dtype] = space.load(0) 44 | result.store( 45 | 0, 46 | (x.load(1) - x.load(0)) / (hu - hd), 47 | ) 48 | 49 | hu = space.load(x.size - 1) 50 | hd = space.load(x.size - 2) 51 | result.store( 52 | x.size - 1, 53 | (x.load(x.size - 1) - x.load(x.size - 2)) / (hu - hd), 54 | ) 55 | 56 | for i in range(1, x.size - 1): 57 | var hu: Scalar[dtype] = space.load(i + 1) - space.load(i) 58 | var hd: Scalar[dtype] = space.load(i) - space.load(i - 1) 59 | var fi: Scalar[dtype] = ( 60 | hd**2 * x.load(i + 1) 61 | + (hu**2 - hd**2) * x.load(i) 62 | - hu**2 * x.load(i - 1) 63 | ) / (hu * hd * (hu + hd)) 64 | result.store(i, fi) 65 | 66 | return result^ 67 | 68 | 69 | # naive loop implementation, optimize later 70 | fn trapz[ 71 | dtype: DType = DType.float64 72 | ](y: NDArray[dtype], x: NDArray[dtype]) raises -> Scalar[dtype]: 73 | """ 74 | Compute the integral of y over x using the trapezoidal rule. 75 | 76 | Parameters: 77 | dtype: The element type. 78 | 79 | Args: 80 | y: An array. 81 | x: An array. 82 | 83 | Constraints: 84 | `x` and `y` must have the same shape. 85 | `fdtype` must be a floating-point type if `idtype` is not a floating-point type. 86 | 87 | Returns: 88 | The integral of y over x using the trapezoidal rule. 89 | """ 90 | constrained[ 91 | is_inttype[dtype]() and not is_floattype[dtype](), 92 | ( 93 | "output dtype `Fdtype` must be a floating-point type if input dtype" 94 | " `Idtype` is not a floating-point type" 95 | ), 96 | ]() 97 | 98 | if x.shape != y.shape: 99 | raise Error("x and y must have the same shape") 100 | 101 | var integral: Scalar[dtype] = 0.0 102 | for i in range(x.size - 1): 103 | var temp = (x.load(i + 1) - x.load(i)) * ( 104 | y.load(i) + y.load(i + 1) 105 | ) / 2.0 106 | integral += temp 107 | return integral 108 | -------------------------------------------------------------------------------- /numojo/routines/math/exponents.mojo: -------------------------------------------------------------------------------- 1 | # ===------------------------------------------------------------------------===# 2 | # Exponents and logarithms 3 | # ===------------------------------------------------------------------------===# 4 | 5 | 6 | import math 7 | from algorithm import parallelize 8 | from algorithm import Static2DTileUnitFunc as Tile2DFunc 9 | from utils import Variant 10 | 11 | import numojo.core._math_funcs as _mf 12 | from numojo.core.ndarray import NDArray 13 | 14 | alias ln = log 15 | """ 16 | Natural Log equivelent to log 17 | """ 18 | 19 | 20 | fn exp[ 21 | dtype: DType, backend: _mf.Backend = _mf.Vectorized 22 | ](array: NDArray[dtype]) raises -> NDArray[dtype]: 23 | """ 24 | Calculate element-wise euler's constant(e) to the power of NDArray[i]. 25 | 26 | Parameters: 27 | dtype: The element type. 28 | backend: Sets utility function origin, defaults to `Vectorized`. 29 | 30 | Args: 31 | array: A NDArray. 32 | 33 | Returns: 34 | A NDArray with the shape of `NDArray` with values equal to the 35 | e to the power of the value in the original NDArray at each position. 36 | """ 37 | return backend().math_func_1_array_in_one_array_out[dtype, math.exp](array) 38 | 39 | 40 | fn exp2[ 41 | dtype: DType, backend: _mf.Backend = _mf.Vectorized 42 | ](array: NDArray[dtype]) raises -> NDArray[dtype]: 43 | """ 44 | Calculate element-wise two to the power of NDArray[i]. 45 | 46 | Parameters: 47 | dtype: The element type. 48 | backend: Sets utility function origin, defaults to `Vectorized`. 49 | 50 | Args: 51 | array: A NDArray. 52 | 53 | Returns: 54 | A NDArray with the shape of `NDArray` with values equal to the 55 | 2 to the power of the value in the original NDArray at each position. 56 | """ 57 | return backend().math_func_1_array_in_one_array_out[dtype, math.exp2](array) 58 | 59 | 60 | fn expm1[ 61 | dtype: DType, backend: _mf.Backend = _mf.Vectorized 62 | ](array: NDArray[dtype]) raises -> NDArray[dtype]: 63 | """ 64 | Calculate element-wise euler's constant(e) to the power of NDArray[i] minus1. 65 | 66 | Parameters: 67 | dtype: The element type. 68 | backend: Sets utility function origin, defaults to `Vectorized`. 69 | 70 | Args: 71 | array: A NDArray. 72 | 73 | Returns: 74 | A NDArray with the shape of `NDArray` with values equal to the negative one plus 75 | e to the power of the value in the original NDArray at each position. 76 | """ 77 | return backend().math_func_1_array_in_one_array_out[dtype, math.expm1]( 78 | array 79 | ) 80 | 81 | 82 | fn log[ 83 | dtype: DType, backend: _mf.Backend = _mf.Vectorized 84 | ](array: NDArray[dtype]) raises -> NDArray[dtype]: 85 | """ 86 | Element-wise natural logarithm of NDArray. 87 | 88 | Parameters: 89 | dtype: The element type. 90 | backend: Sets utility function origin, defaults to `Vectorized`. 91 | 92 | Args: 93 | array: A NDArray. 94 | 95 | Returns: 96 | A NDArray equal to ln(NDArray). 97 | """ 98 | return backend().math_func_1_array_in_one_array_out[dtype, math.log](array) 99 | 100 | 101 | fn log2[ 102 | dtype: DType, backend: _mf.Backend = _mf.Vectorized 103 | ](array: NDArray[dtype]) raises -> NDArray[dtype]: 104 | """ 105 | Element-wise logarithm base two of NDArray. 106 | 107 | Parameters: 108 | dtype: The element type. 109 | backend: Sets utility function origin, defaults to `Vectorized`. 110 | 111 | Args: 112 | array: A NDArray. 113 | 114 | Returns: 115 | A NDArray equal to log_2(NDArray). 116 | """ 117 | return backend().math_func_1_array_in_one_array_out[dtype, math.log2](array) 118 | 119 | 120 | fn log10[ 121 | dtype: DType, backend: _mf.Backend = _mf.Vectorized 122 | ](array: NDArray[dtype]) raises -> NDArray[dtype]: 123 | """ 124 | Element-wise logarithm base ten of NDArray. 125 | 126 | Parameters: 127 | dtype: The element type. 128 | backend: Sets utility function origin, defaults to `Vectorized`. 129 | 130 | Args: 131 | array: A NDArray. 132 | 133 | Returns: 134 | A NDArray equal to log_10(NDArray). 135 | """ 136 | return backend().math_func_1_array_in_one_array_out[dtype, math.log10]( 137 | array 138 | ) 139 | 140 | 141 | fn log1p[ 142 | dtype: DType, backend: _mf.Backend = _mf.Vectorized 143 | ](array: NDArray[dtype]) raises -> NDArray[dtype]: 144 | """ 145 | Element-wise natural logarithm of 1 plus NDArray. 146 | 147 | Parameters: 148 | dtype: The element type. 149 | backend: Sets utility function origin, defaults to `Vectorized`. 150 | 151 | Args: 152 | array: A NDArray. 153 | 154 | Returns: 155 | A NDArray equal to ln(NDArray+1). 156 | """ 157 | return backend().math_func_1_array_in_one_array_out[dtype, math.log1p]( 158 | array 159 | ) 160 | -------------------------------------------------------------------------------- /numojo/routines/math/floating.mojo: -------------------------------------------------------------------------------- 1 | # ===------------------------------------------------------------------------===# 2 | # Floating point routines 3 | # ===------------------------------------------------------------------------===# 4 | 5 | import math 6 | from algorithm import parallelize 7 | from algorithm import Static2DTileUnitFunc as Tile2DFunc 8 | from utils import Variant 9 | 10 | import numojo.core._math_funcs as _mf 11 | from numojo.core.ndarray import NDArray 12 | 13 | 14 | fn copysign[ 15 | dtype: DType, backend: _mf.Backend = _mf.Vectorized 16 | ](array1: NDArray[dtype], array2: NDArray[dtype]) raises -> NDArray[dtype]: 17 | """ 18 | Copy the sign of the first NDArray and apply it to the second NDArray. 19 | 20 | Constraints: 21 | Both arrays must have the same shapes. 22 | 23 | Parameters: 24 | dtype: The element type. 25 | backend: Sets utility function origin, defaults to `Vectorized`. 26 | 27 | Args: 28 | array1: A NDArray. 29 | array2: A NDArray. 30 | 31 | Returns: 32 | The second NDArray multipied by the sign of the first NDArray. 33 | """ 34 | return backend().math_func_2_array_in_one_array_out[dtype, math.copysign]( 35 | array1, array2 36 | ) 37 | -------------------------------------------------------------------------------- /numojo/routines/math/hyper.mojo: -------------------------------------------------------------------------------- 1 | """ 2 | Implements Hyperbolic functions for arrays. 3 | """ 4 | # ===----------------------------------------------------------------------=== # 5 | # Hyperbolic functions 6 | # ===----------------------------------------------------------------------=== # 7 | 8 | import math 9 | 10 | import numojo.core._math_funcs as _mf 11 | from numojo.core.ndarray import NDArray 12 | from numojo.core.matrix import Matrix 13 | import numojo.core.matrix as matrix 14 | 15 | # TODO: add dtype in backends and pass it here. 16 | 17 | # ===------------------------------------------------------------------------===# 18 | # Inverse Hyperbolic Trig 19 | # ===------------------------------------------------------------------------===# 20 | 21 | 22 | fn arccosh[dtype: DType](A: Matrix[dtype]) -> Matrix[dtype]: 23 | return matrix._arithmetic_func_matrix_to_matrix[dtype, math.acosh](A) 24 | 25 | 26 | fn acosh[ 27 | dtype: DType, backend: _mf.Backend = _mf.Vectorized 28 | ](array: NDArray[dtype]) raises -> NDArray[dtype]: 29 | """ 30 | Apply acosh also known as inverse hyperbolic cosine . 31 | 32 | Parameters: 33 | dtype: The element type. 34 | backend: Sets utility function origin, defaults to `Vectorized. 35 | 36 | Args: 37 | array: An Array. 38 | 39 | Returns: 40 | The element-wise acosh of `array` in radians. 41 | """ 42 | return backend().math_func_1_array_in_one_array_out[dtype, math.acosh]( 43 | array 44 | ) 45 | 46 | 47 | fn acosh[dtype: DType](A: Matrix[dtype]) -> Matrix[dtype]: 48 | return matrix._arithmetic_func_matrix_to_matrix[dtype, math.acosh](A) 49 | 50 | 51 | fn arcsinh[dtype: DType](A: Matrix[dtype]) -> Matrix[dtype]: 52 | return matrix._arithmetic_func_matrix_to_matrix[dtype, math.asinh](A) 53 | 54 | 55 | fn asinh[ 56 | dtype: DType, backend: _mf.Backend = _mf.Vectorized 57 | ](array: NDArray[dtype]) raises -> NDArray[dtype]: 58 | """ 59 | Apply asinh also known as inverse hyperbolic sine . 60 | 61 | Parameters: 62 | dtype: The element type. 63 | backend: Sets utility function origin, defaults to `Vectorized. 64 | 65 | Args: 66 | array: An Array. 67 | 68 | Returns: 69 | The element-wise asinh of `array` in radians. 70 | """ 71 | return backend().math_func_1_array_in_one_array_out[dtype, math.asinh]( 72 | array 73 | ) 74 | 75 | 76 | fn asinh[dtype: DType](A: Matrix[dtype]) -> Matrix[dtype]: 77 | return matrix._arithmetic_func_matrix_to_matrix[dtype, math.asinh](A) 78 | 79 | 80 | fn arctanh[dtype: DType](A: Matrix[dtype]) -> Matrix[dtype]: 81 | return matrix._arithmetic_func_matrix_to_matrix[dtype, math.atanh](A) 82 | 83 | 84 | fn atanh[ 85 | dtype: DType, backend: _mf.Backend = _mf.Vectorized 86 | ](array: NDArray[dtype]) raises -> NDArray[dtype]: 87 | """ 88 | Apply atanh also known as inverse hyperbolic tangent . 89 | 90 | Parameters: 91 | dtype: The element type. 92 | backend: Sets utility function origin, defaults to `Vectorized. 93 | 94 | Args: 95 | array: An Array. 96 | 97 | Returns: 98 | The element-wise atanh of `array` in radians. 99 | """ 100 | return backend().math_func_1_array_in_one_array_out[dtype, math.atanh]( 101 | array 102 | ) 103 | 104 | 105 | fn atanh[dtype: DType](A: Matrix[dtype]) -> Matrix[dtype]: 106 | return matrix._arithmetic_func_matrix_to_matrix[dtype, math.atanh](A) 107 | 108 | 109 | # ===------------------------------------------------------------------------===# 110 | # Hyperbolic Trig 111 | # ===------------------------------------------------------------------------===# 112 | 113 | 114 | fn cosh[ 115 | dtype: DType, backend: _mf.Backend = _mf.Vectorized 116 | ](array: NDArray[dtype]) raises -> NDArray[dtype]: 117 | """ 118 | Apply cosh also known as hyperbolic cosine . 119 | 120 | Parameters: 121 | dtype: The element type. 122 | backend: Sets utility function origin, defaults to `Vectorized. 123 | 124 | Args: 125 | array: An Array assumed to be in radian. 126 | 127 | Returns: 128 | The element-wise cosh of `array`. 129 | """ 130 | return backend().math_func_1_array_in_one_array_out[dtype, math.cosh](array) 131 | 132 | 133 | fn cosh[dtype: DType](A: Matrix[dtype]) -> Matrix[dtype]: 134 | return matrix._arithmetic_func_matrix_to_matrix[dtype, math.cosh](A) 135 | 136 | 137 | fn sinh[ 138 | dtype: DType, backend: _mf.Backend = _mf.Vectorized 139 | ](array: NDArray[dtype]) raises -> NDArray[dtype]: 140 | """ 141 | Apply sin also known as hyperbolic sine . 142 | 143 | Parameters: 144 | dtype: The element type. 145 | backend: Sets utility function origin, defaults to `Vectorized. 146 | 147 | Args: 148 | array: An Array assumed to be in radian. 149 | 150 | Returns: 151 | The element-wise sinh of `array`. 152 | """ 153 | return backend().math_func_1_array_in_one_array_out[dtype, math.sinh](array) 154 | 155 | 156 | fn sinh[dtype: DType](A: Matrix[dtype]) -> Matrix[dtype]: 157 | return matrix._arithmetic_func_matrix_to_matrix[dtype, math.sinh](A) 158 | 159 | 160 | fn tanh[ 161 | dtype: DType, backend: _mf.Backend = _mf.Vectorized 162 | ](array: NDArray[dtype]) raises -> NDArray[dtype]: 163 | """ 164 | Apply tan also known as hyperbolic tangent . 165 | 166 | Parameters: 167 | dtype: The element type. 168 | backend: Sets utility function origin, defaults to `Vectorized. 169 | 170 | Args: 171 | array: An Array assumed to be in radian. 172 | 173 | Returns: 174 | The element-wise tanh of `array`. 175 | """ 176 | return backend().math_func_1_array_in_one_array_out[dtype, math.tanh](array) 177 | 178 | 179 | fn tanh[dtype: DType](A: Matrix[dtype]) -> Matrix[dtype]: 180 | return matrix._arithmetic_func_matrix_to_matrix[dtype, math.tanh](A) 181 | -------------------------------------------------------------------------------- /numojo/routines/math/misc.mojo: -------------------------------------------------------------------------------- 1 | # ===----------------------------------------------------------------------=== # 2 | # Distributed under the Apache 2.0 License with LLVM Exceptions. 3 | # See LICENSE and the LLVM License for more information. 4 | # https://github.com/Mojo-Numerics-and-Algorithms-group/NuMojo/blob/main/LICENSE 5 | # https://llvm.org/LICENSE.txt 6 | # ===----------------------------------------------------------------------=== # 7 | 8 | # ===------------------------------------------------------------------------===# 9 | # Miscellaneous mathematical functions 10 | # ===------------------------------------------------------------------------===# 11 | 12 | from algorithm import parallelize, vectorize 13 | from algorithm import Static2DTileUnitFunc as Tile2DFunc 14 | import builtin.math as builtin_math 15 | import stdlib.math.math as stdlib_math 16 | from sys import simdwidthof 17 | from utils import Variant 18 | 19 | import numojo.core._math_funcs as _mf 20 | from numojo.core.ndarray import NDArray 21 | 22 | 23 | fn cbrt[ 24 | dtype: DType, backend: _mf.Backend = _mf.Vectorized 25 | ](array: NDArray[dtype]) raises -> NDArray[dtype]: 26 | """ 27 | Element-wise cuberoot of NDArray. 28 | 29 | Constraints: 30 | Both arrays must have the same shapes. 31 | 32 | Parameters: 33 | dtype: The element type. 34 | backend: Sets utility function origin, defaults to `Vectorized`. 35 | 36 | Args: 37 | array: A NDArray. 38 | 39 | Returns: 40 | A NDArray equal to NDArray**(1/3). 41 | """ 42 | return backend().math_func_1_array_in_one_array_out[ 43 | dtype, stdlib_math.cbrt 44 | ](array) 45 | 46 | 47 | fn clip[ 48 | dtype: DType, // 49 | ](a: NDArray[dtype], a_min: Scalar[dtype], a_max: Scalar[dtype]) -> NDArray[ 50 | dtype 51 | ]: 52 | """ 53 | Limit the values in an array between [a_min, a_max]. 54 | If a_min is greater than a_max, the value is equal to a_max. 55 | 56 | Parameters: 57 | dtype: The data type. 58 | 59 | Args: 60 | a: A array. 61 | a_min: The minimum value. 62 | a_max: The maximum value. 63 | 64 | Returns: 65 | An array with the clipped values. 66 | """ 67 | 68 | var res = a # Deep copy of the array 69 | 70 | for i in range(res.size): 71 | if res._buf.ptr[i] < a_min: 72 | res._buf.ptr[i] = a_min 73 | if res._buf.ptr[i] > a_max: 74 | res._buf.ptr[i] = a_max 75 | 76 | return res 77 | 78 | 79 | fn _mt_rsqrt[ 80 | dtype: DType, simd_width: Int 81 | ](value: SIMD[dtype, simd_width]) -> SIMD[dtype, simd_width]: 82 | """ 83 | Element-wise reciprocal squareroot of SIMD. 84 | Parameters: 85 | dtype: The element type. 86 | simd_width: The SIMD width. 87 | Args: 88 | value: A SIMD vector. 89 | Returns: 90 | A SIMD equal to 1/SIMD**(1/2). 91 | """ 92 | return stdlib_math.sqrt(SIMD.__truediv__(1, value)) 93 | 94 | 95 | fn rsqrt[ 96 | dtype: DType, backend: _mf.Backend = _mf.Vectorized 97 | ](array: NDArray[dtype]) raises -> NDArray[dtype]: 98 | """ 99 | Element-wise reciprocal squareroot of NDArray. 100 | 101 | Parameters: 102 | dtype: The element type. 103 | backend: Sets utility function origin, defaults to `Vectorized`. 104 | 105 | Args: 106 | array: A NDArray. 107 | 108 | Returns: 109 | A NDArray equal to 1/NDArray**(1/2). 110 | """ 111 | return backend().math_func_1_array_in_one_array_out[dtype, _mt_rsqrt](array) 112 | 113 | 114 | fn sqrt[ 115 | dtype: DType, backend: _mf.Backend = _mf.Vectorized 116 | ](array: NDArray[dtype]) raises -> NDArray[dtype]: 117 | """ 118 | Element-wise square root of NDArray. 119 | 120 | Parameters: 121 | dtype: The element type. 122 | backend: Sets utility function origin, defaults to `Vectorized`. 123 | 124 | Args: 125 | array: A NDArray. 126 | 127 | Returns: 128 | A NDArray equal to NDArray**(1/2). 129 | """ 130 | return backend().math_func_1_array_in_one_array_out[ 131 | dtype, stdlib_math.sqrt 132 | ](array) 133 | 134 | 135 | # this is a temporary doc, write a more explanatory one 136 | fn scalb[ 137 | dtype: DType, backend: _mf.Backend = _mf.Vectorized 138 | ](array1: NDArray[dtype], array2: NDArray[dtype]) raises -> NDArray[dtype]: 139 | """ 140 | Calculate the scalb of array1 and array2. 141 | 142 | Parameters: 143 | dtype: The element type. 144 | backend: Sets utility function origin, defaults to `Vectorized`. 145 | 146 | Args: 147 | array1: A NDArray. 148 | array2: A NDArray. 149 | 150 | Returns: 151 | A NDArray with the shape of `NDArray` with values equal to the negative one plus 152 | e to the power of the value in the original NDArray at each position. 153 | """ 154 | return backend().math_func_2_array_in_one_array_out[ 155 | dtype, stdlib_math.scalb 156 | ](array1, array2) 157 | -------------------------------------------------------------------------------- /numojo/routines/math/products.mojo: -------------------------------------------------------------------------------- 1 | from algorithm.functional import parallelize, vectorize 2 | from sys import simdwidthof 3 | 4 | from numojo.core.ndarray import NDArray 5 | import numojo.core.matrix as matrix 6 | from numojo.core.matrix import Matrix 7 | from numojo.routines.creation import ones 8 | 9 | 10 | fn prod[dtype: DType](A: NDArray[dtype]) raises -> Scalar[dtype]: 11 | """ 12 | Returns products of all items in the array. 13 | 14 | Example: 15 | ```console 16 | > print(A) 17 | [[ 0.1315377950668335 0.458650141954422 0.21895918250083923 ] 18 | [ 0.67886471748352051 0.93469291925430298 0.51941639184951782 ] 19 | [ 0.034572109580039978 0.52970021963119507 0.007698186207562685 ]] 20 | 2-D array Shape: [3, 3] DType: float32 21 | 22 | > print(nm.prod(A)) 23 | 6.1377261317829834e-07 24 | ``` 25 | 26 | Args: 27 | A: NDArray. 28 | 29 | Returns: 30 | Scalar. 31 | """ 32 | 33 | alias width: Int = simdwidthof[dtype]() 34 | var res = Scalar[dtype](1) 35 | 36 | @parameter 37 | fn cal_vec[width: Int](i: Int): 38 | res *= A._buf.ptr.load[width=width](i).reduce_mul() 39 | 40 | vectorize[cal_vec, width](A.size) 41 | return res 42 | 43 | 44 | fn prod[ 45 | dtype: DType 46 | ](A: NDArray[dtype], owned axis: Int) raises -> NDArray[dtype]: 47 | """ 48 | Returns products of array elements over a given axis. 49 | 50 | Args: 51 | A: NDArray. 52 | axis: The axis along which the product is performed. 53 | 54 | Returns: 55 | An NDArray. 56 | """ 57 | 58 | if axis < 0: 59 | axis += A.ndim 60 | if (axis < 0) or (axis >= A.ndim): 61 | raise Error( 62 | String("Invalid index: index out of bound [0, {}).").format(A.ndim) 63 | ) 64 | 65 | var result_shape: List[Int] = List[Int]() 66 | var size_of_axis: Int = A.shape[axis] 67 | var slices: List[Slice] = List[Slice]() 68 | for i in range(A.ndim): 69 | if i != axis: 70 | result_shape.append(A.shape[i]) 71 | slices.append(Slice(0, A.shape[i])) 72 | else: 73 | slices.append(Slice(0, 0)) # Temp value 74 | var result = ones[dtype](NDArrayShape(result_shape)) 75 | for i in range(size_of_axis): 76 | slices[axis] = Slice(i, i + 1) 77 | var arr_slice = A[slices] 78 | result *= arr_slice 79 | 80 | return result 81 | 82 | 83 | fn prod[dtype: DType](A: Matrix[dtype]) -> Scalar[dtype]: 84 | """ 85 | Product of all items in the Matrix. 86 | 87 | Args: 88 | A: Matrix. 89 | """ 90 | var res = Scalar[dtype](1) 91 | alias width: Int = simdwidthof[dtype]() 92 | 93 | @parameter 94 | fn cal_vec[width: Int](i: Int): 95 | res = res * A._buf.ptr.load[width=width](i).reduce_mul() 96 | 97 | vectorize[cal_vec, width](A.size) 98 | return res 99 | 100 | 101 | fn prod[dtype: DType](A: Matrix[dtype], axis: Int) raises -> Matrix[dtype]: 102 | """ 103 | Product of items in a Matrix along the axis. 104 | 105 | Args: 106 | A: Matrix. 107 | axis: 0 or 1. 108 | 109 | Example: 110 | ```mojo 111 | from numojo import Matrix 112 | var A = Matrix.rand(shape=(100, 100)) 113 | print(mat.prod(A, axis=0)) 114 | print(mat.prod(A, axis=1)) 115 | ``` 116 | """ 117 | 118 | alias width: Int = simdwidthof[dtype]() 119 | 120 | if axis == 0: 121 | var B = Matrix.ones[dtype](shape=(1, A.shape[1])) 122 | 123 | for i in range(A.shape[0]): 124 | 125 | @parameter 126 | fn cal_vec_sum[width: Int](j: Int): 127 | B._store[width]( 128 | 0, j, B._load[width](0, j) * A._load[width](i, j) 129 | ) 130 | 131 | vectorize[cal_vec_sum, width](A.shape[1]) 132 | 133 | return B^ 134 | 135 | elif axis == 1: 136 | var B = Matrix.ones[dtype](shape=(A.shape[0], 1)) 137 | 138 | @parameter 139 | fn cal_rows(i: Int): 140 | @parameter 141 | fn cal_vec[width: Int](j: Int): 142 | B._store( 143 | i, 144 | 0, 145 | B._load(i, 0) * A._load[width=width](i, j).reduce_mul(), 146 | ) 147 | 148 | vectorize[cal_vec, width](A.shape[1]) 149 | 150 | parallelize[cal_rows](A.shape[0], A.shape[0]) 151 | return B^ 152 | 153 | else: 154 | raise Error(String("The axis can either be 1 or 0!")) 155 | 156 | 157 | fn cumprod[dtype: DType](A: NDArray[dtype]) raises -> NDArray[dtype]: 158 | """ 159 | Returns cumprod of all items of an array. 160 | The array is flattened before cumprod. 161 | 162 | Parameters: 163 | dtype: The element type. 164 | 165 | Args: 166 | A: NDArray. 167 | 168 | Returns: 169 | Cumprod of all items of an array. 170 | """ 171 | 172 | if A.ndim == 1: 173 | var B = A 174 | for i in range(A.size - 1): 175 | B._buf.ptr[i + 1] *= B._buf.ptr[i] 176 | return B^ 177 | 178 | else: 179 | return cumprod(A.flatten(), axis=-1) 180 | 181 | 182 | fn cumprod[ 183 | dtype: DType 184 | ](owned A: NDArray[dtype], owned axis: Int) raises -> NDArray[dtype]: 185 | """ 186 | Returns cumprod of array by axis. 187 | 188 | Parameters: 189 | dtype: The element type. 190 | 191 | Args: 192 | A: NDArray. 193 | axis: Axis. 194 | 195 | Returns: 196 | Cumprod of array by axis. 197 | """ 198 | 199 | if axis < 0: 200 | axis += A.ndim 201 | if (axis < 0) or (axis >= A.ndim): 202 | raise Error( 203 | String("Invalid index: index out of bound [0, {}).").format(A.ndim) 204 | ) 205 | 206 | var I = NDArray[DType.index](Shape(A.size)) 207 | var ptr = I._buf.ptr 208 | 209 | var _shape = A.shape._move_axis_to_end(axis) 210 | var _strides = A.strides._move_axis_to_end(axis) 211 | 212 | numojo.core.utility._traverse_buffer_according_to_shape_and_strides( 213 | ptr, _shape, _strides 214 | ) 215 | 216 | for i in range(0, A.size, A.shape[axis]): 217 | for j in range(A.shape[axis] - 1): 218 | A._buf.ptr[I._buf.ptr[i + j + 1]] *= A._buf.ptr[I._buf.ptr[i + j]] 219 | 220 | return A^ 221 | 222 | 223 | fn cumprod[dtype: DType](owned A: Matrix[dtype]) -> Matrix[dtype]: 224 | """ 225 | Cumprod of flattened matrix. 226 | 227 | Args: 228 | A: Matrix. 229 | 230 | Example: 231 | ```mojo 232 | from numojo import Matrix 233 | var A = Matrix.rand(shape=(100, 100)) 234 | print(mat.cumprod(A)) 235 | ``` 236 | """ 237 | 238 | A.resize(shape=(1, A.size)) 239 | 240 | for i in range(1, A.size): 241 | A._buf.ptr[i] *= A._buf.ptr[i - 1] 242 | 243 | return A^ 244 | 245 | 246 | fn cumprod[ 247 | dtype: DType 248 | ](owned A: Matrix[dtype], axis: Int) raises -> Matrix[dtype]: 249 | """ 250 | Cumprod of Matrix along the axis. 251 | 252 | Args: 253 | A: Matrix. 254 | axis: 0 or 1. 255 | 256 | Example: 257 | ```mojo 258 | from numojo import Matrix 259 | var A = Matrix.rand(shape=(100, 100)) 260 | print(mat.cumprod(A, axis=0)) 261 | print(mat.cumprod(A, axis=1)) 262 | ``` 263 | """ 264 | 265 | alias width: Int = simdwidthof[dtype]() 266 | 267 | if axis == 0: 268 | for i in range(1, A.shape[0]): 269 | 270 | @parameter 271 | fn cal_vec[width: Int](j: Int): 272 | A._store[width]( 273 | i, j, A._load[width](i - 1, j) * A._load[width](i, j) 274 | ) 275 | 276 | vectorize[cal_vec, width](A.shape[1]) 277 | 278 | return A^ 279 | 280 | elif axis == 1: 281 | return transpose(cumprod(transpose(A), axis=0)) 282 | 283 | else: 284 | raise Error(String("The axis can either be 1 or 0!")) 285 | -------------------------------------------------------------------------------- /numojo/routines/math/rounding.mojo: -------------------------------------------------------------------------------- 1 | # ===------------------------------------------------------------------------===# 2 | # Rounding 3 | # ===------------------------------------------------------------------------===# 4 | 5 | from builtin import math as builtin_math 6 | from algorithm import parallelize 7 | from algorithm import Static2DTileUnitFunc as Tile2DFunc 8 | from utils import Variant 9 | from utils.numerics import nextafter as builtin_nextafter 10 | 11 | import numojo.core._math_funcs as _mf 12 | from numojo.core.ndarray import NDArray 13 | import numojo.core.matrix as matrix 14 | from numojo.core.matrix import Matrix 15 | 16 | 17 | fn round[ 18 | dtype: DType 19 | ](owned A: Matrix[dtype], decimals: Int = 0) -> Matrix[dtype]: 20 | # FIXME 21 | # The built-in `round` function is not working now. 22 | # It will be fixed in future. 23 | 24 | for i in range(A.size): 25 | A._buf.ptr[i] = builtin_math.round(A._buf.ptr[i], ndigits=decimals) 26 | 27 | return A^ 28 | 29 | 30 | fn tabs[ 31 | dtype: DType, backend: _mf.Backend = _mf.Vectorized 32 | ](array: NDArray[dtype]) raises -> NDArray[dtype]: 33 | """ 34 | Element-wise absolute value of NDArray. 35 | 36 | Parameters: 37 | dtype: The element type. 38 | backend: Sets utility function origin, defaults to `Vectorized`. 39 | 40 | Args: 41 | array: A NDArray. 42 | 43 | Returns: 44 | A NDArray equal to abs(NDArray). 45 | """ 46 | return backend().math_func_1_array_in_one_array_out[dtype, SIMD.__abs__]( 47 | array 48 | ) 49 | 50 | 51 | fn tfloor[ 52 | dtype: DType, backend: _mf.Backend = _mf.Vectorized 53 | ](array: NDArray[dtype]) raises -> NDArray[dtype]: 54 | """ 55 | Element-wise round down to nearest whole number of NDArray. 56 | 57 | Parameters: 58 | dtype: The element type. 59 | backend: Sets utility function origin, defaults to `Vectorized`. 60 | 61 | Args: 62 | array: A NDArray. 63 | 64 | Returns: 65 | A NDArray equal to floor(NDArray). 66 | """ 67 | return backend().math_func_1_array_in_one_array_out[dtype, SIMD.__floor__]( 68 | array 69 | ) 70 | 71 | 72 | fn tceil[ 73 | dtype: DType, backend: _mf.Backend = _mf.Vectorized 74 | ](array: NDArray[dtype]) raises -> NDArray[dtype]: 75 | """ 76 | Element-wise round up to nearest whole number of NDArray. 77 | 78 | Parameters: 79 | dtype: The element type. 80 | backend: Sets utility function origin, defaults to `Vectorized`. 81 | 82 | Args: 83 | array: A NDArray. 84 | 85 | Returns: 86 | A NDArray equal to ceil(NDArray). 87 | """ 88 | return backend().math_func_1_array_in_one_array_out[dtype, SIMD.__ceil__]( 89 | array 90 | ) 91 | 92 | 93 | fn ttrunc[ 94 | dtype: DType, backend: _mf.Backend = _mf.Vectorized 95 | ](array: NDArray[dtype]) raises -> NDArray[dtype]: 96 | """ 97 | Element-wise remove decimal value from float whole number of NDArray. 98 | 99 | Parameters: 100 | dtype: The element type. 101 | backend: Sets utility function origin, defaults to `Vectorized`. 102 | 103 | Args: 104 | array: A NDArray. 105 | 106 | Returns: 107 | A NDArray equal to trunc(NDArray). 108 | """ 109 | return backend().math_func_1_array_in_one_array_out[dtype, SIMD.__trunc__]( 110 | array 111 | ) 112 | 113 | 114 | fn tround[ 115 | dtype: DType, backend: _mf.Backend = _mf.Vectorized 116 | ](array: NDArray[dtype]) raises -> NDArray[dtype]: 117 | """ 118 | Element-wise round NDArray to whole number. 119 | 120 | Parameters: 121 | dtype: The element type. 122 | backend: Sets utility function origin, defaults to `Vectorized`. 123 | 124 | Args: 125 | array: A NDArray. 126 | 127 | Returns: 128 | A NDArray equal to trunc(NDArray). 129 | """ 130 | return backend().math_func_1_array_in_one_array_out[dtype, SIMD.__round__]( 131 | array 132 | ) 133 | 134 | 135 | fn roundeven[ 136 | dtype: DType, backend: _mf.Backend = _mf.Vectorized 137 | ](array: NDArray[dtype]) raises -> NDArray[dtype]: 138 | """ 139 | Performs element-wise banker's rounding on the elements of a NDArray. 140 | 141 | Parameters: 142 | dtype: The dtype of the input and output array. 143 | backend: Sets utility function origin, defaults to `Vectorized`. 144 | 145 | Args: 146 | array: Array to perform rounding on. 147 | 148 | Returns: 149 | The element-wise banker's rounding of NDArray. 150 | 151 | This rounding goes to the nearest integer with ties toward the nearest even integer. 152 | """ 153 | return backend().math_func_1_array_in_one_array_out[dtype, SIMD.roundeven]( 154 | array 155 | ) 156 | 157 | 158 | # fn round_half_down[ 159 | # dtype: DType, backend: _mf.Backend = _mf.Vectorized 160 | # ](NDArray: NDArray[dtype]) -> NDArray[dtype]: 161 | # """ 162 | # Rounds ties towards the smaller integer. 163 | 164 | # Parameters: 165 | # dtype: The dtype of the input and output array. 166 | # backend: Sets utility function origin, defaults to `Vectorized`. 167 | 168 | # Args: 169 | # NDArray: array to perform rounding on. 170 | 171 | # Returns: 172 | # The element-wise rounding of x evaluating ties towards the smaller integer. 173 | # """ 174 | # return backend().math_func_1_array_in_one_array_out[ 175 | # dtype, SIMD.__round_half_down 176 | # ](NDArray) 177 | 178 | 179 | # fn round_half_up[ 180 | # dtype: DType, backend: _mf.Backend = _mf.Vectorized 181 | # ](NDArray: NDArray[dtype]) -> NDArray[dtype]: 182 | # """ 183 | # Rounds ties towards the larger integer. 184 | 185 | # Parameters: 186 | # dtype: The dtype of the input and output array. 187 | # backend: Sets utility function origin, defaults to `Vectorized`. 188 | 189 | # Args: 190 | # NDArray: array to perform rounding on. 191 | 192 | # Returns: 193 | # The element-wise rounding of x evaluating ties towards the larger integer. 194 | # """ 195 | # return backend().math_func_1_array_in_one_array_out[ 196 | # dtype, math.round_half_up 197 | # ](NDArray) 198 | 199 | 200 | fn nextafter[ 201 | dtype: DType, backend: _mf.Backend = _mf.Vectorized 202 | ](array1: NDArray[dtype], array2: NDArray[dtype]) raises -> NDArray[dtype]: 203 | """ 204 | Computes the nextafter of the inputs. 205 | 206 | Parameters: 207 | dtype: The dtype of the input and output array. Constraints: must be a floating-point type. 208 | backend: Sets utility function origin, default to `Vectorized`. 209 | 210 | 211 | Args: 212 | array1: The first input argument. 213 | array2: The second input argument. 214 | 215 | Returns: 216 | The nextafter of the inputs. 217 | """ 218 | return backend().math_func_2_array_in_one_array_out[ 219 | dtype, builtin_nextafter 220 | ](array1, array2) 221 | -------------------------------------------------------------------------------- /numojo/routines/math/sums.mojo: -------------------------------------------------------------------------------- 1 | from sys import simdwidthof 2 | from algorithm import parallelize, vectorize 3 | 4 | from numojo.core.ndarray import NDArray 5 | from numojo.core.matrix import Matrix 6 | from numojo.routines.creation import zeros 7 | 8 | 9 | fn sum[dtype: DType](A: NDArray[dtype]) -> Scalar[dtype]: 10 | """ 11 | Returns sum of all items in the array. 12 | 13 | Example: 14 | ```console 15 | > print(A) 16 | [[ 0.1315377950668335 0.458650141954422 0.21895918250083923 ] 17 | [ 0.67886471748352051 0.93469291925430298 0.51941639184951782 ] 18 | [ 0.034572109580039978 0.52970021963119507 0.007698186207562685 ]] 19 | 2-D array Shape: [3, 3] DType: float32 20 | > print(nm.sum(A)) 21 | 3.5140917301177979 22 | ``` 23 | 24 | Args: 25 | A: NDArray. 26 | 27 | Returns: 28 | Scalar. 29 | """ 30 | 31 | alias width: Int = simdwidthof[dtype]() 32 | var res = Scalar[dtype](0) 33 | 34 | @parameter 35 | fn cal_vec[width: Int](i: Int): 36 | res += A._buf.ptr.load[width=width](i).reduce_add() 37 | 38 | vectorize[cal_vec, width](A.size) 39 | return res 40 | 41 | 42 | fn sum[dtype: DType](A: NDArray[dtype], axis: Int) raises -> NDArray[dtype]: 43 | """ 44 | Returns sums of array elements over a given axis. 45 | 46 | Example: 47 | ```mojo 48 | import numojo as nm 49 | var A = nm.random.randn(100, 100) 50 | print(nm.sum(A, axis=0)) 51 | ``` 52 | 53 | Raises: 54 | Error: If the axis is out of bound. 55 | Error: If the number of dimensions is 1. 56 | 57 | Args: 58 | A: NDArray. 59 | axis: The axis along which the sum is performed. 60 | 61 | Returns: 62 | An NDArray. 63 | """ 64 | 65 | var normalized_axis = axis 66 | if normalized_axis < 0: 67 | normalized_axis += A.ndim 68 | 69 | if (normalized_axis < 0) or (normalized_axis >= A.ndim): 70 | raise Error( 71 | String("Axis {} out of bound [0, {}).").format(axis, A.ndim) 72 | ) 73 | if A.ndim == 1: 74 | raise Error( 75 | String( 76 | "`numojo.routines.math.sums.sum()`: " 77 | "Cannot sum over axis for 1-d array. " 78 | "Please remove the `axis` argument." 79 | ) 80 | ) 81 | 82 | var result_shape: List[Int] = List[Int]() 83 | var size_of_axis: Int = A.shape[normalized_axis] 84 | var slices: List[Slice] = List[Slice]() 85 | for i in range(A.ndim): 86 | if i != normalized_axis: 87 | result_shape.append(A.shape[i]) 88 | slices.append(Slice(0, A.shape[i])) 89 | else: 90 | slices.append(Slice(0, 0)) # Temp value 91 | var result = zeros[dtype](NDArrayShape(result_shape)) 92 | for i in range(size_of_axis): 93 | slices[normalized_axis] = Slice(i, i + 1) 94 | var arr_slice = A[slices] 95 | result += arr_slice 96 | 97 | return result 98 | 99 | 100 | fn sum[dtype: DType](A: Matrix[dtype]) -> Scalar[dtype]: 101 | """ 102 | Sum up all items in the Matrix. 103 | 104 | Args: 105 | A: Matrix. 106 | 107 | Example: 108 | ```mojo 109 | from numojo import Matrix 110 | var A = Matrix.rand(shape=(100, 100)) 111 | print(mat.sum(A)) 112 | ``` 113 | """ 114 | var res = Scalar[dtype](0) 115 | alias width: Int = simdwidthof[dtype]() 116 | 117 | @parameter 118 | fn cal_vec[width: Int](i: Int): 119 | res = res + A._buf.ptr.load[width=width](i).reduce_add() 120 | 121 | vectorize[cal_vec, width](A.size) 122 | return res 123 | 124 | 125 | fn sum[dtype: DType](A: Matrix[dtype], axis: Int) raises -> Matrix[dtype]: 126 | """ 127 | Sum up the items in a Matrix along the axis. 128 | 129 | Args: 130 | A: Matrix. 131 | axis: 0 or 1. 132 | 133 | Example: 134 | ```mojo 135 | from numojo import Matrix 136 | var A = Matrix.rand(shape=(100, 100)) 137 | print(mat.sum(A, axis=0)) 138 | print(mat.sum(A, axis=1)) 139 | ``` 140 | """ 141 | 142 | alias width: Int = simdwidthof[dtype]() 143 | 144 | if axis == 0: 145 | var B = Matrix.zeros[dtype](shape=(1, A.shape[1])) 146 | 147 | for i in range(A.shape[0]): 148 | 149 | @parameter 150 | fn cal_vec_sum[width: Int](j: Int): 151 | B._store[width]( 152 | 0, j, B._load[width](0, j) + A._load[width](i, j) 153 | ) 154 | 155 | vectorize[cal_vec_sum, width](A.shape[1]) 156 | 157 | return B^ 158 | 159 | elif axis == 1: 160 | var B = Matrix.zeros[dtype](shape=(A.shape[0], 1)) 161 | 162 | @parameter 163 | fn cal_rows(i: Int): 164 | @parameter 165 | fn cal_vec[width: Int](j: Int): 166 | B._store( 167 | i, 168 | 0, 169 | B._load(i, 0) + A._load[width=width](i, j).reduce_add(), 170 | ) 171 | 172 | vectorize[cal_vec, width](A.shape[1]) 173 | 174 | parallelize[cal_rows](A.shape[0], A.shape[0]) 175 | return B^ 176 | 177 | else: 178 | raise Error(String("The axis can either be 1 or 0!")) 179 | 180 | 181 | fn cumsum[dtype: DType](A: NDArray[dtype]) raises -> NDArray[dtype]: 182 | """ 183 | Returns cumsum of all items of an array. 184 | The array is flattened before cumsum. 185 | 186 | Parameters: 187 | dtype: The element type. 188 | 189 | Args: 190 | A: NDArray. 191 | 192 | Returns: 193 | Cumsum of all items of an array. 194 | """ 195 | 196 | if A.ndim == 1: 197 | var B = A 198 | for i in range(A.size - 1): 199 | B._buf.ptr[i + 1] += B._buf.ptr[i] 200 | return B^ 201 | 202 | else: 203 | return cumsum(A.flatten(), axis=-1) 204 | 205 | 206 | fn cumsum[ 207 | dtype: DType 208 | ](owned A: NDArray[dtype], owned axis: Int) raises -> NDArray[dtype]: 209 | """ 210 | Returns cumsum of array by axis. 211 | 212 | Parameters: 213 | dtype: The element type. 214 | 215 | Args: 216 | A: NDArray. 217 | axis: Axis. 218 | 219 | Returns: 220 | Cumsum of array by axis. 221 | """ 222 | 223 | if axis < 0: 224 | axis += A.ndim 225 | if (axis < 0) or (axis >= A.ndim): 226 | raise Error( 227 | String("Invalid index: index out of bound [0, {}).").format(A.ndim) 228 | ) 229 | 230 | var I = NDArray[DType.index](Shape(A.size)) 231 | var ptr = I._buf.ptr 232 | 233 | var _shape = A.shape._move_axis_to_end(axis) 234 | var _strides = A.strides._move_axis_to_end(axis) 235 | 236 | numojo.core.utility._traverse_buffer_according_to_shape_and_strides( 237 | ptr, _shape, _strides 238 | ) 239 | 240 | for i in range(0, A.size, A.shape[axis]): 241 | for j in range(A.shape[axis] - 1): 242 | A._buf.ptr[Int(I._buf.ptr[i + j + 1])] += A._buf.ptr[ 243 | Int(I._buf.ptr[i + j]) 244 | ] 245 | 246 | return A^ 247 | 248 | 249 | fn cumsum[dtype: DType](owned A: Matrix[dtype]) -> Matrix[dtype]: 250 | """ 251 | Cumsum of flattened matrix. 252 | 253 | Args: 254 | A: Matrix. 255 | 256 | Example: 257 | ```mojo 258 | from numojo import Matrix 259 | var A = Matrix.rand(shape=(100, 100)) 260 | print(mat.cumsum(A)) 261 | ``` 262 | """ 263 | 264 | A.resize(shape=(1, A.size)) 265 | 266 | for i in range(1, A.size): 267 | A._buf.ptr[i] += A._buf.ptr[i - 1] 268 | 269 | return A^ 270 | 271 | 272 | fn cumsum[ 273 | dtype: DType 274 | ](owned A: Matrix[dtype], axis: Int) raises -> Matrix[dtype]: 275 | """ 276 | Cumsum of Matrix along the axis. 277 | 278 | Args: 279 | A: Matrix. 280 | axis: 0 or 1. 281 | 282 | Example: 283 | ```mojo 284 | from numojo import Matrix 285 | var A = Matrix.rand(shape=(100, 100)) 286 | print(mat.cumsum(A, axis=0)) 287 | print(mat.cumsum(A, axis=1)) 288 | ``` 289 | """ 290 | 291 | alias width: Int = simdwidthof[dtype]() 292 | 293 | if axis == 0: 294 | for i in range(1, A.shape[0]): 295 | 296 | @parameter 297 | fn cal_vec_sum[width: Int](j: Int): 298 | A._store[width]( 299 | i, j, A._load[width](i - 1, j) + A._load[width](i, j) 300 | ) 301 | 302 | vectorize[cal_vec_sum, width](A.shape[1]) 303 | 304 | return A^ 305 | 306 | elif axis == 1: 307 | return transpose(cumsum(transpose(A), axis=0)) 308 | 309 | else: 310 | raise Error(String("The axis can either be 1 or 0!")) 311 | -------------------------------------------------------------------------------- /numojo/routines/math/trig.mojo: -------------------------------------------------------------------------------- 1 | """ 2 | Implements Trigonometry functions for arrays. 3 | """ 4 | # ===----------------------------------------------------------------------=== # 5 | # Trigonometric functions 6 | # ===----------------------------------------------------------------------=== # 7 | 8 | import math 9 | 10 | import numojo.core._math_funcs as _mf 11 | from numojo.core.ndarray import NDArray 12 | from numojo.core.matrix import Matrix 13 | import numojo.core.matrix as matrix 14 | 15 | from numojo.routines.math.misc import sqrt 16 | from numojo.routines.math.arithmetic import fma 17 | 18 | # TODO: add dtype in backends and pass it here. 19 | 20 | # ===------------------------------------------------------------------------===# 21 | # Inverse Trig 22 | # ===------------------------------------------------------------------------===# 23 | 24 | 25 | fn arccos[dtype: DType](A: Matrix[dtype]) -> Matrix[dtype]: 26 | return matrix._arithmetic_func_matrix_to_matrix[dtype, math.acos](A) 27 | 28 | 29 | fn acos[ 30 | dtype: DType, backend: _mf.Backend = _mf.Vectorized 31 | ](array: NDArray[dtype]) raises -> NDArray[dtype]: 32 | """ 33 | Apply acos also known as inverse cosine . 34 | 35 | Parameters: 36 | dtype: The element type. 37 | backend: Sets utility function origin, defaults to `Vectorized. 38 | 39 | Args: 40 | array: An Array. 41 | 42 | Returns: 43 | The element-wise acos of `array` in radians. 44 | """ 45 | return backend().math_func_1_array_in_one_array_out[dtype, math.acos](array) 46 | 47 | 48 | fn acos[dtype: DType](A: Matrix[dtype]) -> Matrix[dtype]: 49 | return matrix._arithmetic_func_matrix_to_matrix[dtype, math.acos](A) 50 | 51 | 52 | fn arcsin[dtype: DType](A: Matrix[dtype]) -> Matrix[dtype]: 53 | return matrix._arithmetic_func_matrix_to_matrix[dtype, math.asin](A) 54 | 55 | 56 | fn asin[ 57 | dtype: DType, backend: _mf.Backend = _mf.Vectorized 58 | ](array: NDArray[dtype]) raises -> NDArray[dtype]: 59 | """ 60 | Apply asin also known as inverse sine . 61 | 62 | Parameters: 63 | dtype: The element type. 64 | backend: Sets utility function origin, defaults to `Vectorized. 65 | 66 | Args: 67 | array: An Array. 68 | 69 | Returns: 70 | The element-wise asin of `array` in radians. 71 | """ 72 | return backend().math_func_1_array_in_one_array_out[dtype, math.asin](array) 73 | 74 | 75 | fn asin[dtype: DType](A: Matrix[dtype]) -> Matrix[dtype]: 76 | return matrix._arithmetic_func_matrix_to_matrix[dtype, math.asin](A) 77 | 78 | 79 | fn arctan[dtype: DType](A: Matrix[dtype]) -> Matrix[dtype]: 80 | return matrix._arithmetic_func_matrix_to_matrix[dtype, math.atan](A) 81 | 82 | 83 | fn atan[ 84 | dtype: DType, backend: _mf.Backend = _mf.Vectorized 85 | ](array: NDArray[dtype]) raises -> NDArray[dtype]: 86 | """ 87 | Apply atan also known as inverse tangent . 88 | 89 | Parameters: 90 | dtype: The element type. 91 | backend: Sets utility function origin, defaults to `Vectorized. 92 | 93 | Args: 94 | array: An Array. 95 | 96 | Returns: 97 | The element-wise atan of `array` in radians. 98 | """ 99 | return backend().math_func_1_array_in_one_array_out[dtype, math.atan](array) 100 | 101 | 102 | fn atan[dtype: DType](A: Matrix[dtype]) -> Matrix[dtype]: 103 | return matrix._arithmetic_func_matrix_to_matrix[dtype, math.atan](A) 104 | 105 | 106 | fn atan2[ 107 | dtype: DType, backend: _mf.Backend = _mf.Vectorized 108 | ](array1: NDArray[dtype], array2: NDArray[dtype]) raises -> NDArray[dtype]: 109 | """ 110 | Apply atan2 also known as inverse tangent. 111 | [atan2 wikipedia](https://en.wikipedia.org/wiki/Atan2). 112 | 113 | Constraints: 114 | Both arrays must have the same shapes. 115 | 116 | Parameters: 117 | dtype: The element type. 118 | backend: Sets utility function origin, defaults to `Vectorized. 119 | 120 | Args: 121 | array1: An Array. 122 | array2: An Array. 123 | 124 | Returns: 125 | The element-wise atan2 of `array1` and`array2` in radians. 126 | """ 127 | return backend().math_func_2_array_in_one_array_out[dtype, math.atan2]( 128 | array1, array2 129 | ) 130 | 131 | 132 | # ===------------------------------------------------------------------------===# 133 | # Trig 134 | # ===------------------------------------------------------------------------===# 135 | 136 | 137 | fn cos[ 138 | dtype: DType, backend: _mf.Backend = _mf.Vectorized 139 | ](array: NDArray[dtype]) raises -> NDArray[dtype]: 140 | """ 141 | Apply cos also known as cosine. 142 | 143 | Parameters: 144 | dtype: The element type. 145 | backend: Sets utility function origin, defaults to `Vectorized. 146 | 147 | Args: 148 | array: An Array assumed to be in radian. 149 | 150 | Returns: 151 | The element-wise cos of `array`. 152 | """ 153 | return backend().math_func_1_array_in_one_array_out[dtype, math.cos](array) 154 | 155 | 156 | fn cos[dtype: DType](A: Matrix[dtype]) -> Matrix[dtype]: 157 | return matrix._arithmetic_func_matrix_to_matrix[dtype, math.cos](A) 158 | 159 | 160 | fn sin[ 161 | dtype: DType, backend: _mf.Backend = _mf.Vectorized 162 | ](array: NDArray[dtype]) raises -> NDArray[dtype]: 163 | """ 164 | Apply sin also known as sine . 165 | 166 | Parameters: 167 | dtype: The element type. 168 | backend: Sets utility function origin, defaults to `Vectorized. 169 | 170 | Args: 171 | array: An Array assumed to be in radian. 172 | 173 | Returns: 174 | The element-wise sin of `array`. 175 | """ 176 | return backend().math_func_1_array_in_one_array_out[dtype, math.sin](array) 177 | 178 | 179 | fn sin[dtype: DType](A: Matrix[dtype]) -> Matrix[dtype]: 180 | return matrix._arithmetic_func_matrix_to_matrix[dtype, math.sin](A) 181 | 182 | 183 | fn tan[ 184 | dtype: DType, backend: _mf.Backend = _mf.Vectorized 185 | ](array: NDArray[dtype]) raises -> NDArray[dtype]: 186 | """ 187 | Apply tan also known as tangent . 188 | 189 | Parameters: 190 | dtype: The element type. 191 | backend: Sets utility function origin, defaults to `Vectorized. 192 | 193 | Args: 194 | array: An Array assumed to be in radian. 195 | 196 | Returns: 197 | The element-wise tan of `array`. 198 | """ 199 | return backend().math_func_1_array_in_one_array_out[dtype, math.tan](array) 200 | 201 | 202 | fn tan[dtype: DType](A: Matrix[dtype]) -> Matrix[dtype]: 203 | return matrix._arithmetic_func_matrix_to_matrix[dtype, math.tan](A) 204 | 205 | 206 | fn hypot[ 207 | dtype: DType, backend: _mf.Backend = _mf.Vectorized 208 | ](array1: NDArray[dtype], array2: NDArray[dtype]) raises -> NDArray[dtype]: 209 | """ 210 | Apply hypot also known as hypotenuse which finds the longest section of a right triangle 211 | given the other two sides. 212 | 213 | Constraints: 214 | Both arrays must have the same shapes. 215 | 216 | Parameters: 217 | dtype: The element type. 218 | backend: Sets utility function origin, defaults to `Vectorized. 219 | 220 | Args: 221 | array1: An Array. 222 | array2: An Array. 223 | 224 | Returns: 225 | The element-wise hypotenuse of `array1` and`array2`. 226 | """ 227 | return backend().math_func_2_array_in_one_array_out[dtype, math.hypot]( 228 | array1, array2 229 | ) 230 | 231 | 232 | fn hypot_fma[ 233 | dtype: DType, backend: _mf.Backend = _mf.Vectorized 234 | ](array1: NDArray[dtype], array2: NDArray[dtype]) raises -> NDArray[dtype]: 235 | """ 236 | Apply hypot also known as hypotenuse which finds the longest section of a right triangle 237 | given the other two sides. 238 | 239 | Constraints: 240 | Both arrays must have the same shapes. 241 | 242 | Parameters: 243 | dtype: The element type. 244 | backend: Sets utility function origin, defaults to `Vectorized. 245 | 246 | Args: 247 | array1: An Array. 248 | array2: An Array. 249 | 250 | Returns: 251 | The element-wise hypotenuse of `array1` and`array2`. 252 | """ 253 | 254 | var array2_squared = fma[dtype, backend=backend]( 255 | array2, array2, SIMD[dtype, 1](0) 256 | ) 257 | return sqrt[dtype, backend=backend]( 258 | fma[dtype, backend=backend](array1, array1, array2_squared) 259 | ) 260 | -------------------------------------------------------------------------------- /numojo/routines/searching.mojo: -------------------------------------------------------------------------------- 1 | # ===----------------------------------------------------------------------=== # 2 | # Searching 3 | # ===----------------------------------------------------------------------=== # 4 | 5 | import math 6 | from algorithm import vectorize 7 | from sys import simdwidthof 8 | from collections.optional import Optional 9 | 10 | from numojo.core.ndarray import NDArray 11 | from numojo.core.ndshape import NDArrayShape 12 | import numojo.core.matrix as matrix 13 | from numojo.core.matrix import Matrix 14 | from numojo.core.utility import is_inttype, is_floattype 15 | from numojo.routines.sorting import binary_sort 16 | from numojo.routines.math.extrema import _max, _min 17 | 18 | 19 | # * for loop version works fine for argmax and argmin, need to vectorize it 20 | fn argmax[dtype: DType](array: NDArray[dtype]) raises -> Int: 21 | """ 22 | Argmax of a array. 23 | 24 | Parameters: 25 | dtype: The element type. 26 | 27 | Args: 28 | array: A array. 29 | Returns: 30 | The index of the maximum value of the array. 31 | """ 32 | if array.size == 0: 33 | raise Error("array is empty") 34 | 35 | var idx: Int = 0 36 | var max_val: Scalar[dtype] = array.load(0) 37 | for i in range(1, array.size): 38 | if array.load(i) > max_val: 39 | max_val = array.load(i) 40 | idx = i 41 | return idx 42 | 43 | 44 | fn argmax[dtype: DType](A: Matrix[dtype]) raises -> Scalar[DType.index]: 45 | """ 46 | Index of the max. It is first flattened before sorting. 47 | """ 48 | 49 | var max_index: Scalar[DType.index] 50 | _, max_index = _max(A, 0, A.size - 1) 51 | 52 | return max_index 53 | 54 | 55 | fn argmax[ 56 | dtype: DType 57 | ](A: Matrix[dtype], axis: Int) raises -> Matrix[DType.index]: 58 | """ 59 | Index of the max along the given axis. 60 | """ 61 | if axis == 1: 62 | var B = Matrix[DType.index](shape=(A.shape[0], 1)) 63 | for i in range(A.shape[0]): 64 | B._store( 65 | i, 66 | 0, 67 | _max(A, start=i * A.strides[0], end=(i + 1) * A.strides[0] - 1)[ 68 | 1 69 | ] 70 | - i * A.strides[0], 71 | ) 72 | return B^ 73 | elif axis == 0: 74 | return transpose(argmax(transpose(A), axis=1)) 75 | else: 76 | raise Error(String("The axis can either be 1 or 0!")) 77 | 78 | 79 | fn argmin[dtype: DType](array: NDArray[dtype]) raises -> Int: 80 | """ 81 | Argmin of a array. 82 | Parameters: 83 | dtype: The element type. 84 | 85 | Args: 86 | array: A array. 87 | Returns: 88 | The index of the minimum value of the array. 89 | """ 90 | if array.size == 0: 91 | raise Error("array is empty") 92 | 93 | var idx: Int = 0 94 | var min_val: Scalar[dtype] = array.load(0) 95 | 96 | for i in range(1, array.size): 97 | if array.load(i) < min_val: 98 | min_val = array.load(i) 99 | idx = i 100 | return idx 101 | 102 | 103 | fn argmin[dtype: DType](A: Matrix[dtype]) raises -> Scalar[DType.index]: 104 | """ 105 | Index of the min. It is first flattened before sorting. 106 | """ 107 | 108 | var min_index: Scalar[DType.index] 109 | _, min_index = _min(A, 0, A.size - 1) 110 | 111 | return min_index 112 | 113 | 114 | fn argmin[ 115 | dtype: DType 116 | ](A: Matrix[dtype], axis: Int) raises -> Matrix[DType.index]: 117 | """ 118 | Index of the min along the given axis. 119 | """ 120 | if axis == 1: 121 | var B = Matrix[DType.index](shape=(A.shape[0], 1)) 122 | for i in range(A.shape[0]): 123 | B._store( 124 | i, 125 | 0, 126 | _min(A, start=i * A.strides[0], end=(i + 1) * A.strides[0] - 1)[ 127 | 1 128 | ] 129 | - i * A.strides[0], 130 | ) 131 | return B^ 132 | elif axis == 0: 133 | return transpose(argmin(transpose(A), axis=1)) 134 | else: 135 | raise Error(String("The axis can either be 1 or 0!")) 136 | -------------------------------------------------------------------------------- /numojo/routines/statistics/__init__.mojo: -------------------------------------------------------------------------------- 1 | from .averages import ( 2 | mean, 3 | max, 4 | min, 5 | mode, 6 | median, 7 | variance, 8 | std, 9 | ) 10 | -------------------------------------------------------------------------------- /numojo/science/__init__.mojo: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Mojo-Numerics-and-Algorithms-group/NuMojo/f0dea371bfd15df65e917e013e11d6177bfc975a/numojo/science/__init__.mojo -------------------------------------------------------------------------------- /numojo/science/signal.mojo: -------------------------------------------------------------------------------- 1 | """ 2 | Implements signal processing. 3 | 4 | It is like `scipy.signal` in Python. 5 | """ 6 | 7 | from numojo.core.ndarray import NDArray 8 | from numojo.core.ndshape import Shape 9 | from numojo.core.item import Item 10 | from numojo.routines.creation import fromstring, zeros 11 | from numojo.routines.math.sums import sum 12 | 13 | 14 | fn convolve2d[ 15 | dtype: DType, //, 16 | ](in1: NDArray[dtype], in2: NDArray[dtype]) raises -> NDArray[dtype]: 17 | """Convolve two 2-dimensional arrays. 18 | 19 | Args: 20 | in1: Input array 1. 21 | in2: Input array 2. It should be of a smaller size of in1. 22 | 23 | Currently, the mode is "valid". 24 | 25 | TODO: Add more modes. 26 | 27 | Example: 28 | ```mojo 29 | import numojo as nm 30 | fn main() raises: 31 | var in1 = nm.routines.creation.fromstring("[[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16]]") 32 | var in2 = nm.routines.creation.fromstring("[[1, 0], [0, -1]]") 33 | print(nm.science.signal.convolve2d(in1, in2)) 34 | ``` 35 | """ 36 | 37 | var in2_mirrored = in2 38 | var length = in2.size 39 | for i in range(length): 40 | in2_mirrored._buf.ptr[i] = in2._buf.ptr[length - i - 1] 41 | 42 | var in1_height = in1.shape[0] 43 | var in1_width = in1.shape[1] 44 | var in2_height = in2_mirrored.shape[0] 45 | var in2_width = in2_mirrored.shape[1] 46 | 47 | var output_height = in1_height - in2_height + 1 48 | var output_width = in1_width - in2_width + 1 49 | 50 | var output = zeros[dtype](Shape(output_height, output_width)) 51 | 52 | for i in range(output_height): 53 | for j in range(output_width): 54 | output[Item(i, j)] = sum( 55 | in1[i : i + in2_height, j : j + in2_width] * in2_mirrored 56 | ) 57 | 58 | return output 59 | -------------------------------------------------------------------------------- /tests/README.md: -------------------------------------------------------------------------------- 1 | # Testing guide 2 | 3 | Inorder to ensure that all of our code works and that a new code doesn't break old code we need to instate a testing policy (which we can hopefully implement a CI/CD process for running). 4 | 5 | Due to the similiarity with NumPy, a Numpy equivelent comparison will be our method for confirming validity. 6 | 7 | The functions contained in utils_for_test.mojo `check` and `check_is_close` will be used to compare NuMojo results with NumPy equivelents. 8 | 9 | Each test must be in it's own def function and begin with the word test example `test_arange`. This allows the `mojo test` command to find it. A single function can cover multiple similiar tests but they should have unique strings to identify which check failed. 10 | -------------------------------------------------------------------------------- /tests/core/test_array_indexing_and_slicing.mojo: -------------------------------------------------------------------------------- 1 | import numojo as nm 2 | from numojo.prelude import * 3 | from testing.testing import assert_true, assert_almost_equal, assert_equal 4 | from utils_for_test import check, check_is_close 5 | from python import Python 6 | 7 | 8 | def test_getitem_scalar(): 9 | var np = Python.import_module("numpy") 10 | 11 | var A = nm.arange(8) 12 | assert_true(A.load(0) == 0, msg=String("`get` fails")) 13 | 14 | 15 | def test_setitem(): 16 | var np = Python.import_module("numpy") 17 | var arr = nm.NDArray(Shape(4, 4)) 18 | var np_arr = arr.to_numpy() 19 | arr.itemset(List(2, 2), 1000) 20 | np_arr[(2, 2)] = 1000 21 | check_is_close(arr, np_arr, "Itemset is broken") 22 | 23 | 24 | def test_slicing_getter1(): 25 | var np = Python.import_module("numpy") 26 | 27 | # Test C-order array slicing 28 | nm_arr = nm.arange[nm.f32](0.0, 24.0, step=1) 29 | nm_arr = nm_arr.reshape(Shape(2, 3, 4), order="C") 30 | np_arr = np.arange(0, 24, dtype=np.float32).reshape(2, 3, 4) 31 | 32 | # Test case 1: Slicing all dimensions 33 | nm_slice1 = nm_arr[:, :, 1:2] 34 | np_sliced1 = np.take( 35 | np.take( 36 | np.take(np_arr, np.arange(0, 2), axis=0), np.arange(0, 3), axis=1 37 | ), 38 | np.arange(1, 2), 39 | axis=2, 40 | ) 41 | np_sliced1 = np.squeeze(np_sliced1, axis=2) 42 | check(nm_slice1, np_sliced1, "3D array slicing (C-order) [:, :, 1:2]") 43 | 44 | 45 | def test_slicing_getter2(): 46 | var np = Python.import_module("numpy") 47 | 48 | # Test C-order array slicing 49 | nm_arr = nm.arange[nm.f32](0.0, 24.0, step=1) 50 | nm_arr = nm_arr.reshape(Shape(2, 3, 4), order="C") 51 | np_arr = np.arange(0, 24, dtype=np.float32).reshape(2, 3, 4) 52 | 53 | # Test case 2: Slicing with start and end indices 54 | nm_slice2 = nm_arr[0:1, 1:3, 2:4] 55 | np_sliced2 = np.take( 56 | np.take( 57 | np.take(np_arr, np.arange(0, 1), axis=0), np.arange(1, 3), axis=1 58 | ), 59 | np.arange(2, 4), 60 | axis=2, 61 | ) 62 | check(nm_slice2, np_sliced2, "3D array slicing (C-order) [0:1, 1:3, 2:4]") 63 | 64 | 65 | def test_slicing_getter3(): 66 | var np = Python.import_module("numpy") 67 | 68 | # Test C-order array slicing 69 | nm_arr = nm.arange[nm.f32](0.0, 24.0, step=1) 70 | nm_arr = nm_arr.reshape(Shape(2, 3, 4), order="C") 71 | np_arr = np.arange(0, 24, dtype=np.float32).reshape(2, 3, 4) 72 | 73 | # Test case 3: Slicing with mixed start, end, and step values 74 | nm_slice3 = nm_arr[1:, 0:2, ::2] 75 | np_sliced3 = np.take( 76 | np.take( 77 | np.take(np_arr, np.arange(1, np_arr.shape[0]), axis=0), 78 | np.arange(0, 2), 79 | axis=1, 80 | ), 81 | np.arange(0, np_arr.shape[2], 2), 82 | axis=2, 83 | ) 84 | check(nm_slice3, np_sliced3, "3D array slicing (C-order) [1:, 0:2, ::2]") 85 | 86 | 87 | def test_slicing_getter4(): 88 | var np = Python.import_module("numpy") 89 | 90 | # Test C-order array slicing 91 | nm_arr = nm.arange[nm.f32](0.0, 24.0, step=1) 92 | nm_arr = nm_arr.reshape(Shape(2, 3, 4), order="C") 93 | np_arr = np.arange(0, 24, dtype=np.float32).reshape(2, 3, 4) 94 | 95 | # Test case 4: Slicing with step 96 | nm_slice4 = nm_arr[::2, ::2, ::2] 97 | np_sliced4 = np.take( 98 | np.take( 99 | np.take(np_arr, np.arange(0, np_arr.shape[0], 2), axis=0), 100 | np.arange(0, np_arr.shape[1], 2), 101 | axis=1, 102 | ), 103 | np.arange(0, np_arr.shape[2], 2), 104 | axis=2, 105 | ) 106 | check(nm_slice4, np_sliced4, "3D array slicing (C-order) [::2, ::2, ::2]") 107 | 108 | 109 | def test_slicing_getter5(): 110 | var np = Python.import_module("numpy") 111 | 112 | # Test C-order array slicing 113 | nm_arr = nm.arange[nm.f32](0.0, 24.0, step=1) 114 | nm_arr = nm_arr.reshape(Shape(2, 3, 4), order="C") 115 | np_arr = np.arange(0, 24, dtype=np.float32).reshape(2, 3, 4) 116 | 117 | # Test case 5: Slicing with combination of integer and slices 118 | nm_slice5 = nm_arr[1:2, :, 1:3] 119 | np_sliced5 = np.take( 120 | np.take(np_arr[1], np.arange(0, np_arr.shape[1]), axis=0), 121 | np.arange(1, 3), 122 | axis=1, 123 | ) 124 | check(nm_slice5, np_sliced5, "3D array slicing (C-order) [1, :, 1:3]") 125 | 126 | 127 | def test_slicing_getter6(): 128 | var np = Python.import_module("numpy") 129 | 130 | var b = nm.arange[i8](60).reshape(Shape(3, 4, 5)) 131 | var ind = nm.array[isize]("[[2,0,1], [1,0,1]]") 132 | var mask = nm.array[boolean]("[1,0,1]") 133 | 134 | var bnp = b.to_numpy() 135 | var indnp = ind.to_numpy() 136 | var masknp = mask.to_numpy() 137 | 138 | check(b[ind], bnp[indnp], "Get by indices array fails") 139 | check(b[mask], bnp[masknp], "Get by mask array fails") 140 | 141 | 142 | # def test_slicing_setter1(): 143 | # var np = Python.import_module("numpy") 144 | 145 | # # Test C-order array slicing 146 | # nm_arr = nm.arange[nm.f32](0.0, 24.0, step=1) 147 | # nm_arr.reshape(2, 3, 4, order="C") 148 | # np_arr = np.arange(0, 24, dtype=np.float32).reshape(2, 3, 4) 149 | 150 | # # Test case 2: Setting a slice with another array 151 | # nm_set_arr = nm.full[nm.f32](2, 2, fill_value=50.0) 152 | # np_set_arr = np.full((1, 2, 2), 50, dtype=np.float32) 153 | # nm_arr[1:2, 1:3, 2:4] = nm_set_arr 154 | # np.put(np_arr, np.ravel_multi_index((np.arange(1, 2), np.arange(1, 3), np.arange(2, 4)), np_arr.shape), np_set_arr.flatten()) 155 | # check(nm_arr, np_arr, "3D array slice setting (C-order) [1:2, 1:3, 2:4] = array") 156 | -------------------------------------------------------------------------------- /tests/core/test_array_methods.mojo: -------------------------------------------------------------------------------- 1 | from python import Python 2 | 3 | from numojo.prelude import * 4 | from testing.testing import assert_true, assert_almost_equal, assert_equal 5 | from utils_for_test import check, check_is_close, check_values_close 6 | 7 | 8 | def test_constructors(): 9 | # Test NDArray constructor with different input types 10 | var arr1 = NDArray[f32](Shape(3, 4, 5)) 11 | assert_true(arr1.ndim == 3, "NDArray constructor: ndim") 12 | assert_true(arr1.shape[0] == 3, "NDArray constructor: shape element 0") 13 | assert_true(arr1.shape[1] == 4, "NDArray constructor: shape element 1") 14 | assert_true(arr1.shape[2] == 5, "NDArray constructor: shape element 2") 15 | assert_true(arr1.size == 60, "NDArray constructor: size") 16 | assert_true(arr1.dtype == DType.float32, "NDArray constructor: dtype") 17 | 18 | var arr2 = NDArray[f32](VariadicList[Int](3, 4, 5)) 19 | assert_true( 20 | arr2.shape[0] == 3, 21 | "NDArray constructor with VariadicList: shape element 0", 22 | ) 23 | assert_true( 24 | arr2.shape[1] == 4, 25 | "NDArray constructor with VariadicList: shape element 1", 26 | ) 27 | assert_true( 28 | arr2.shape[2] == 5, 29 | "NDArray constructor with VariadicList: shape element 2", 30 | ) 31 | 32 | var arr3 = nm.full[f32](Shape(3, 4, 5), fill_value=Scalar[f32](10.0)) 33 | # maybe it's better to return a scalar for arr[0, 0, 0] 34 | assert_equal( 35 | arr3[Item(0, 0, 0)], 10.0, "NDArray constructor with fill value" 36 | ) 37 | 38 | var arr4 = NDArray[f32](List[Int](3, 4, 5)) 39 | assert_true( 40 | arr4.shape[0] == 3, "NDArray constructor with List: shape element 0" 41 | ) 42 | assert_true( 43 | arr4.shape[1] == 4, "NDArray constructor with List: shape element 1" 44 | ) 45 | assert_true( 46 | arr4.shape[2] == 5, "NDArray constructor with List: shape element 2" 47 | ) 48 | 49 | var arr5 = NDArray[f32](NDArrayShape(3, 4, 5)) 50 | assert_true( 51 | arr5.shape[0] == 3, 52 | "NDArray constructor with NDArrayShape: shape element 0", 53 | ) 54 | assert_true( 55 | arr5.shape[1] == 4, 56 | "NDArray constructor with NDArrayShape: shape element 1", 57 | ) 58 | assert_true( 59 | arr5.shape[2] == 5, 60 | "NDArray constructor with NDArrayShape: shape element 2", 61 | ) 62 | 63 | 64 | def test_iterator(): 65 | var py = Python.import_module("builtins") 66 | var np = Python.import_module("numpy") 67 | 68 | var a = nm.arange[i8](24).reshape(Shape(2, 3, 4)) 69 | var anp = a.to_numpy() 70 | var f = nm.arange[i8](24).reshape(Shape(2, 3, 4), order="F") 71 | var fnp = f.to_numpy() 72 | 73 | # NDAxisIter 74 | var a_iter_along_axis = a.iter_along_axis[forward=False](axis=0) 75 | var b = a_iter_along_axis.__next__() == nm.array[i8]("[11, 23]") 76 | assert_true( 77 | b.item(0) == True, 78 | "`_NDAxisIter` breaks", 79 | ) 80 | assert_true( 81 | b.item(1) == True, 82 | "`_NDAxisIter` breaks", 83 | ) 84 | 85 | # NDArrayIter 86 | var a_iter_over_dimension = a.__iter__() 87 | var anp_iter_over_dimension = anp.__iter__() 88 | for i in range(a.shape[0]): 89 | check( 90 | a_iter_over_dimension.__next__(), 91 | anp_iter_over_dimension.__next__(), 92 | "`_NDArrayIter` or `__iter__()` breaks", 93 | ) 94 | 95 | var a_iter_over_dimension_reversed = a.__reversed__() 96 | var anp_iter_over_dimension_reversed = py.reversed(anp) 97 | for i in range(a.shape[0]): 98 | check( 99 | a_iter_over_dimension_reversed.__next__(), 100 | anp_iter_over_dimension_reversed.__next__(), 101 | "`_NDArrayIter` or `__reversed__()` breaks", 102 | ) 103 | 104 | # NDIter of C-order array 105 | var a_nditer = a.nditer() 106 | var anp_nditer = np.nditer(anp) 107 | for i in range(a.size): 108 | check_values_close( 109 | a_nditer.__next__(), 110 | anp_nditer.__next__(), 111 | "`_NDIter` or `nditer()` of C array by order C breaks", 112 | ) 113 | 114 | # NDIter of C-order array 115 | a_nditer = a.nditer() 116 | anp_nditer = np.nditer(anp) 117 | for i in range(a.size): 118 | check_values_close( 119 | a_nditer.ith(i), 120 | anp_nditer.__next__(), 121 | "`_NDIter.ith()` of C array by order C breaks", 122 | ) 123 | 124 | var a_nditer_f = a.nditer(order="F") 125 | var anp_nditer_f = np.nditer(anp, order="F") 126 | for i in range(a.size): 127 | check_values_close( 128 | a_nditer_f.__next__(), 129 | anp_nditer_f.__next__(), 130 | "`_NDIter` or `nditer()` of C array by order F breaks", 131 | ) 132 | 133 | # NDIter of F-order array 134 | var f_nditer = f.nditer() 135 | var fnp_nditer = np.nditer(fnp) 136 | for i in range(f.size): 137 | check_values_close( 138 | f_nditer.__next__(), 139 | fnp_nditer.__next__(), 140 | "`_NDIter` or `nditer()` of F array by order C breaks", 141 | ) 142 | 143 | var f_nditer_f = f.nditer(order="F") 144 | var fnp_nditer_f = np.nditer(fnp, order="F") 145 | for i in range(f.size): 146 | check_values_close( 147 | f_nditer_f.__next__(), 148 | fnp_nditer_f.__next__(), 149 | "`_NDIter` or `nditer()` of F array by order F breaks", 150 | ) 151 | -------------------------------------------------------------------------------- /tests/core/test_bool_masks.mojo: -------------------------------------------------------------------------------- 1 | import numojo as nm 2 | from numojo import * 3 | from testing.testing import assert_true, assert_almost_equal, assert_equal 4 | from utils_for_test import check, check 5 | from python import Python 6 | 7 | 8 | # TODO: there's something wrong with bool comparision even though result looks same. 9 | def test_bool_masks_gt(): 10 | var np = Python.import_module("numpy") 11 | var np_A = np.arange(0, 24, dtype=np.int16).reshape((3, 2, 4)) 12 | var A = nm.arange[nm.i16](0, 24) 13 | A = A.reshape(Shape(3, 2, 4)) 14 | 15 | var np_gt = np_A > 10 16 | var gt = A > Scalar[nm.i16](10) 17 | check(gt, np_gt, "Greater than mask") 18 | 19 | # Test greater than or equal 20 | var np_ge = np_A >= 10 21 | var ge = A >= Scalar[nm.i16](10) 22 | check(ge, np_ge, "Greater than or equal mask") 23 | 24 | 25 | def test_bool_masks_lt(): 26 | var np = Python.import_module("numpy") 27 | 28 | # Create NumPy and NuMojo arrays using arange and reshape 29 | var np_A = np.arange(0, 24, dtype=np.int16).reshape((3, 2, 4)) 30 | var A = nm.arange[nm.i16](0, 24) 31 | A = A.reshape(Shape(3, 2, 4)) 32 | 33 | # Test less than 34 | var np_lt = np_A < 10 35 | var lt = A < Scalar[nm.i16](10) 36 | check(lt, np_lt, "Less than mask") 37 | 38 | # Test less than or equal 39 | var np_le = np_A <= 10 40 | var le = A <= Scalar[nm.i16](10) 41 | check(le, np_le, "Less than or equal mask") 42 | 43 | 44 | def test_bool_masks_eq(): 45 | var np = Python.import_module("numpy") 46 | 47 | # Create NumPy and NuMojo arrays using arange and reshape 48 | var np_A = np.arange(0, 24, dtype=np.int16).reshape((3, 2, 4)) 49 | var A = nm.arange[nm.i16](0, 24) 50 | A = A.reshape(Shape(3, 2, 4)) 51 | 52 | # Test equal 53 | # var np_eq = np_A == 10 # has a python interop problem. 54 | var np_eq = np.equal(np_A, 10) 55 | var eq = A == Scalar[nm.i16](10) 56 | check(eq, np_eq, "Equal mask") 57 | 58 | # Test not equal 59 | # var np_ne = np_A != 10 # has a python interop problem. 60 | var np_ne = np.not_equal(np_A, 10) 61 | var ne = A != Scalar[nm.i16](10) 62 | check(ne, np_ne, "Not equal mask") 63 | 64 | # Test masked array 65 | var np_mask = np_A[np_A > 10] 66 | var mask = A[A > Scalar[nm.i16](10)] 67 | check(mask, np_mask, "Masked array") 68 | -------------------------------------------------------------------------------- /tests/core/test_complexArray.mojo: -------------------------------------------------------------------------------- 1 | from testing import assert_equal, assert_almost_equal 2 | from numojo import * 3 | 4 | 5 | fn test_complex_array_init() raises: 6 | """Test initialization of ComplexArray.""" 7 | var c1 = ComplexNDArray[cf32](Shape(2, 2)) 8 | c1.itemset(0, ComplexSIMD[cf32](1.0, 2.0)) 9 | c1.itemset(1, ComplexSIMD[cf32](3.0, 4.0)) 10 | c1.itemset(2, ComplexSIMD[cf32](5.0, 6.0)) 11 | c1.itemset(3, ComplexSIMD[cf32](7.0, 8.0)) 12 | assert_almost_equal(c1.item(0).re, 1.0, "init failed") 13 | assert_almost_equal(c1.item(0).im, 2.0, "init failed") 14 | 15 | 16 | fn test_complex_array_add() raises: 17 | """Test addition of ComplexArray numbers.""" 18 | var c1 = ComplexNDArray[cf32](Shape(2, 2)) 19 | var c2 = ComplexNDArray[cf32](Shape(2, 2)) 20 | c1.itemset(0, ComplexSIMD[cf32](1.0, 2.0)) 21 | c1.itemset(1, ComplexSIMD[cf32](3.0, 4.0)) 22 | c1.itemset(2, ComplexSIMD[cf32](5.0, 6.0)) 23 | c1.itemset(3, ComplexSIMD[cf32](7.0, 8.0)) 24 | c2.itemset(0, ComplexSIMD[cf32](1.0, 2.0)) 25 | c2.itemset(1, ComplexSIMD[cf32](3.0, 4.0)) 26 | c2.itemset(2, ComplexSIMD[cf32](5.0, 6.0)) 27 | c2.itemset(3, ComplexSIMD[cf32](7.0, 8.0)) 28 | 29 | var sum = c1 + c2 30 | 31 | assert_almost_equal(sum.item(0).re, 2.0, "add failed") 32 | assert_almost_equal(sum.item(0).im, 4.0, "add failed") 33 | assert_almost_equal(sum.item(1).re, 6.0, "add failed") 34 | assert_almost_equal(sum.item(1).im, 8.0, "add failed") 35 | assert_almost_equal(sum.item(2).re, 10.0, "add failed") 36 | assert_almost_equal(sum.item(2).im, 12.0, "add failed") 37 | assert_almost_equal(sum.item(3).re, 14.0, "add failed") 38 | assert_almost_equal(sum.item(3).im, 16.0, "add failed") 39 | 40 | 41 | fn test_complex_array_sub() raises: 42 | """Test subtraction of ComplexArray numbers.""" 43 | var c1 = ComplexNDArray[cf32](Shape(2, 2)) 44 | var c2 = ComplexNDArray[cf32](Shape(2, 2)) 45 | c1.itemset(0, ComplexSIMD[cf32](1.0, 2.0)) 46 | c1.itemset(1, ComplexSIMD[cf32](3.0, 4.0)) 47 | c1.itemset(2, ComplexSIMD[cf32](5.0, 6.0)) 48 | c1.itemset(3, ComplexSIMD[cf32](7.0, 8.0)) 49 | 50 | c2.itemset(0, ComplexSIMD[cf32](3.0, 4.0)) 51 | c2.itemset(1, ComplexSIMD[cf32](5.0, 6.0)) 52 | c2.itemset(2, ComplexSIMD[cf32](7.0, 8.0)) 53 | c2.itemset(3, ComplexSIMD[cf32](9.0, 10.0)) 54 | 55 | var diff = c1 - c2 56 | 57 | assert_almost_equal(diff.item(0).re, -2.0, "sub failed") 58 | assert_almost_equal(diff.item(0).im, -2.0, "sub failed") 59 | assert_almost_equal(diff.item(1).re, -2.0, "sub failed") 60 | assert_almost_equal(diff.item(1).im, -2.0, "sub failed") 61 | assert_almost_equal(diff.item(2).re, -2.0, "sub failed") 62 | assert_almost_equal(diff.item(2).im, -2.0, "sub failed") 63 | assert_almost_equal(diff.item(3).re, -2.0, "sub failed") 64 | assert_almost_equal(diff.item(3).im, -2.0, "sub failed") 65 | 66 | 67 | fn test_complex_array_mul() raises: 68 | """Test multiplication of ComplexArray numbers.""" 69 | var c1 = ComplexNDArray[cf32](Shape(2, 2)) 70 | var c2 = ComplexNDArray[cf32](Shape(2, 2)) 71 | c1.itemset(0, ComplexSIMD[cf32](1.0, 2.0)) 72 | c1.itemset(1, ComplexSIMD[cf32](3.0, 4.0)) 73 | c1.itemset(2, ComplexSIMD[cf32](5.0, 6.0)) 74 | c1.itemset(3, ComplexSIMD[cf32](7.0, 8.0)) 75 | 76 | c2.itemset(0, ComplexSIMD[cf32](1.0, 2.0)) 77 | c2.itemset(1, ComplexSIMD[cf32](3.0, 4.0)) 78 | c2.itemset(2, ComplexSIMD[cf32](5.0, 6.0)) 79 | c2.itemset(3, ComplexSIMD[cf32](7.0, 8.0)) 80 | 81 | var prod = c1 * c2 82 | 83 | assert_almost_equal(prod.item(0).re, -3.0, "mul failed") 84 | assert_almost_equal(prod.item(0).im, 4.0, "mul failed") 85 | 86 | 87 | fn test_complex_array_div() raises: 88 | """Test division of ComplexArray numbers.""" 89 | var c1 = ComplexNDArray[cf32](Shape(2, 2)) 90 | var c2 = ComplexNDArray[cf32](Shape(2, 2)) 91 | c1.itemset(0, ComplexSIMD[cf32](1.0, 2.0)) 92 | c1.itemset(1, ComplexSIMD[cf32](3.0, 4.0)) 93 | c1.itemset(2, ComplexSIMD[cf32](5.0, 6.0)) 94 | c1.itemset(3, ComplexSIMD[cf32](7.0, 8.0)) 95 | 96 | c2.itemset(0, ComplexSIMD[cf32](3.0, 4.0)) 97 | c2.itemset(1, ComplexSIMD[cf32](5.0, 6.0)) 98 | c2.itemset(2, ComplexSIMD[cf32](7.0, 8.0)) 99 | c2.itemset(3, ComplexSIMD[cf32](9.0, 10.0)) 100 | 101 | var quot = c1 / c2 102 | 103 | assert_almost_equal(quot.item(0).re, 0.44, "div failed") 104 | assert_almost_equal(quot.item(0).im, 0.08, "div failed") 105 | -------------------------------------------------------------------------------- /tests/core/test_complexSIMD.mojo: -------------------------------------------------------------------------------- 1 | from testing import assert_equal, assert_almost_equal 2 | from numojo import * 3 | 4 | 5 | fn test_complex_init() raises: 6 | """Test initialization of ComplexSIMD.""" 7 | var c1 = ComplexSIMD[cf32](1.0, 2.0) 8 | assert_equal(c1.re, 1.0, "init failed") 9 | assert_equal(c1.im, 2.0, "init failed") 10 | 11 | var c2 = ComplexSIMD[cf32](c1) 12 | assert_equal(c2.re, c1.re) 13 | assert_equal(c2.im, c1.im) 14 | 15 | 16 | fn test_complex_add() raises: 17 | """Test addition of ComplexSIMD numbers.""" 18 | var c1 = ComplexSIMD[cf32](1.0, 2.0) 19 | var c2 = ComplexSIMD[cf32](3.0, 4.0) 20 | 21 | var sum = c1 + c2 22 | assert_equal(sum.re, 4.0, "addition failed") 23 | assert_equal(sum.im, 6.0, "addition failed") 24 | 25 | var c3 = c1 26 | c3 += c2 27 | assert_equal(c3.re, 4.0, "addition failed") 28 | assert_equal(c3.im, 6.0, "addition failed") 29 | 30 | 31 | fn test_complex_sub() raises: 32 | """Test subtraction of ComplexSIMD numbers.""" 33 | var c1 = ComplexSIMD[cf32](3.0, 4.0) 34 | var c2 = ComplexSIMD[cf32](1.0, 2.0) 35 | 36 | var diff = c1 - c2 37 | assert_equal(diff.re, 2.0, "subtraction failed") 38 | assert_equal(diff.im, 2.0, "subtraction failed") 39 | 40 | var c3 = c1 41 | c3 -= c2 42 | assert_equal(c3.re, 2.0, "subtraction failed") 43 | assert_equal(c3.im, 2.0, "subtraction failed") 44 | 45 | 46 | fn test_complex_mul() raises: 47 | """Test multiplication of ComplexSIMD numbers.""" 48 | var c1 = ComplexSIMD[cf32](1.0, 2.0) 49 | var c2 = ComplexSIMD[cf32](3.0, 4.0) 50 | 51 | # (1 + 2i)(3 + 4i) = (1*3 - 2*4) + (1*4 + 2*3)i = -5 + 10i 52 | var prod = c1 * c2 53 | assert_equal(prod.re, -5.0, "multiplication failed") 54 | assert_equal(prod.im, 10.0, "multiplication failed") 55 | 56 | var c3 = c1 57 | c3 *= c2 58 | assert_equal(c3.re, -5.0, "multiplication failed") 59 | assert_equal(c3.im, 10.0, "multiplication failed") 60 | 61 | 62 | fn test_complex_div() raises: 63 | """Test division of ComplexSIMD numbers.""" 64 | var c1 = ComplexSIMD[cf32](1.0, 2.0) 65 | var c2 = ComplexSIMD[cf32](3.0, 4.0) 66 | 67 | # (1 + 2i)/(3 + 4i) = (1*3 + 2*4 + (2*3 - 1*4)i)/(3^2 + 4^2) 68 | # = (11 + 2i)/25 69 | var quot = c1 / c2 70 | assert_almost_equal(quot.re, 0.44, " division failed") 71 | assert_almost_equal(quot.im, 0.08, " division failed") 72 | -------------------------------------------------------------------------------- /tests/core/test_shape_strides_item.mojo: -------------------------------------------------------------------------------- 1 | from numojo.prelude import * 2 | from testing.testing import assert_true, assert_almost_equal, assert_equal 3 | from utils_for_test import check, check_is_close 4 | 5 | 6 | def test_shape(): 7 | var A = nm.NDArrayShape(2, 3, 4) 8 | assert_true( 9 | A[-1] == 4, 10 | msg=String("`NDArrayShape.__getitem__()` fails: may overflow"), 11 | ) 12 | 13 | 14 | def test_strides(): 15 | var A = nm.NDArrayStrides(2, 3, 4) 16 | assert_true( 17 | A[-1] == 4, 18 | msg=String("`NDArrayStrides.__getitem__()` fails: may overflow"), 19 | ) 20 | assert_true( 21 | A[-2] == 3, 22 | msg=String("`NDArrayStrides.__getitem__()` fails: may overflow"), 23 | ) 24 | 25 | 26 | def test_item(): 27 | var A = nm.Item(2, 3, 4) 28 | assert_true( 29 | A[-1] == 4, 30 | msg=String("`NDArrayStrides.__getitem__()` fails: may overflow"), 31 | ) 32 | -------------------------------------------------------------------------------- /tests/routines/test_functional.mojo: -------------------------------------------------------------------------------- 1 | # ===----------------------------------------------------------------------=== # 2 | # Distributed under the Apache 2.0 License with LLVM Exceptions. 3 | # See LICENSE and the LLVM License for more information. 4 | # https://github.com/Mojo-Numerics-and-Algorithms-group/NuMojo/blob/main/LICENSE 5 | # https://llvm.org/LICENSE.txt 6 | # ===----------------------------------------------------------------------=== # 7 | """ 8 | Test functional programming module `numojo.routines.functional`. 9 | """ 10 | from python import Python 11 | from testing.testing import assert_true, assert_almost_equal, assert_equal 12 | from utils_for_test import check, check_is_close 13 | 14 | from numojo.prelude import * 15 | 16 | 17 | fn test_apply_along_axis() raises: 18 | var np = Python.import_module("numpy") 19 | var a = nm.random.randn(Shape(4, 8, 16)) 20 | var anp = a.to_numpy() 21 | var b = nm.reshape(a, a.shape, order="F") 22 | var bnp = b.to_numpy() 23 | 24 | for i in range(a.ndim): 25 | check( 26 | nm.apply_along_axis[nm.sorting.quick_sort_1d](a, axis=i), 27 | np.apply_along_axis(np.sort, axis=i, arr=anp), 28 | String( 29 | "`apply_along_axis` C-order array along axis {} is broken" 30 | ).format(i), 31 | ) 32 | check( 33 | nm.apply_along_axis[nm.sorting.quick_sort_1d](b, axis=i), 34 | np.apply_along_axis(np.sort, axis=i, arr=bnp), 35 | String( 36 | "`apply_along_axis` F-order array along axis {} is broken" 37 | ).format(i), 38 | ) 39 | -------------------------------------------------------------------------------- /tests/routines/test_linalg.mojo: -------------------------------------------------------------------------------- 1 | import numojo as nm 2 | from numojo.prelude import * 3 | from python import Python, PythonObject 4 | from utils_for_test import check, check_is_close, check_values_close 5 | 6 | # ===-----------------------------------------------------------------------===# 7 | # Matmul 8 | # ===-----------------------------------------------------------------------===# 9 | # ! MATMUL RESULTS IN A SEGMENTATION FAULT EXCEPT FOR NAIVE ONE, BUT NAIVE OUTPUTS WRONG VALUES 10 | 11 | 12 | def test_matmul_small(): 13 | var np = Python.import_module("numpy") 14 | var arr = nm.ones[i8](Shape(4, 4)) 15 | var np_arr = np.ones((4, 4), dtype=np.int8) 16 | check_is_close( 17 | arr @ arr, np.matmul(np_arr, np_arr), "Dunder matmul is broken" 18 | ) 19 | 20 | 21 | def test_matmul(): 22 | var np = Python.import_module("numpy") 23 | var arr = nm.arange[nm.f64](0, 100) 24 | arr.resize(Shape(10, 10)) 25 | var np_arr = np.arange(0, 100).reshape(10, 10) 26 | check_is_close( 27 | arr @ arr, np.matmul(np_arr, np_arr), "Dunder matmul is broken" 28 | ) 29 | # The only matmul that currently works is par (__matmul__) 30 | # check_is_close(nm.matmul_tiled_unrolled_parallelized(arr,arr),np.matmul(np_arr,np_arr),"TUP matmul is broken") 31 | 32 | 33 | def test_matmul_4dx4d(): 34 | var np = Python.import_module("numpy") 35 | var A = nm.random.randn(2, 3, 4, 5) 36 | var B = nm.random.randn(2, 3, 5, 4) 37 | check_is_close( 38 | A @ B, 39 | np.matmul(A.to_numpy(), B.to_numpy()), 40 | "`matmul_4dx4d` is broken", 41 | ) 42 | 43 | 44 | def test_matmul_8dx8d(): 45 | var np = Python.import_module("numpy") 46 | var A = nm.random.randn(2, 3, 4, 5, 6, 5, 4, 3) 47 | var B = nm.random.randn(2, 3, 4, 5, 6, 5, 3, 2) 48 | check_is_close( 49 | A @ B, 50 | np.matmul(A.to_numpy(), B.to_numpy()), 51 | "`matmul_8dx8d` is broken", 52 | ) 53 | 54 | 55 | def test_matmul_1dx2d(): 56 | var np = Python.import_module("numpy") 57 | var arr1 = nm.random.randn(4) 58 | var arr2 = nm.random.randn(4, 8) 59 | var nparr1 = arr1.to_numpy() 60 | var nparr2 = arr2.to_numpy() 61 | check_is_close( 62 | arr1 @ arr2, np.matmul(nparr1, nparr2), "Dunder matmul is broken" 63 | ) 64 | 65 | 66 | def test_matmul_2dx1d(): 67 | var np = Python.import_module("numpy") 68 | var arr1 = nm.random.randn(11, 4) 69 | var arr2 = nm.random.randn(4) 70 | var nparr1 = arr1.to_numpy() 71 | var nparr2 = arr2.to_numpy() 72 | check_is_close( 73 | arr1 @ arr2, np.matmul(nparr1, nparr2), "Dunder matmul is broken" 74 | ) 75 | 76 | 77 | # ! The `inv` is broken, it outputs -INF for some values 78 | def test_inv(): 79 | var np = Python.import_module("numpy") 80 | var arr = nm.core.random.rand(100, 100) 81 | var np_arr = arr.to_numpy() 82 | check_is_close( 83 | nm.math.linalg.inv(arr), np.linalg.inv(np_arr), "Inverse is broken" 84 | ) 85 | 86 | 87 | # ! The `solve` is broken, it outputs -INF, nan, 0 etc for some values 88 | def test_solve(): 89 | var np = Python.import_module("numpy") 90 | var A = nm.core.random.randn(100, 100) 91 | var B = nm.core.random.randn(100, 50) 92 | var A_np = A.to_numpy() 93 | var B_np = B.to_numpy() 94 | check_is_close( 95 | nm.linalg.solve(A, B), 96 | np.linalg.solve(A_np, B_np), 97 | "Solve is broken", 98 | ) 99 | 100 | 101 | def norms(): 102 | var np = Python.import_module("numpy") 103 | var arr = nm.core.random.rand(20, 20) 104 | var np_arr = arr.to_numpy() 105 | check_values_close( 106 | nm.math.linalg.det(arr), np.linalg.det(np_arr), "`det` is broken" 107 | ) 108 | 109 | 110 | def test_misc(): 111 | var np = Python.import_module("numpy") 112 | var arr = nm.core.random.rand(4, 8) 113 | var np_arr = arr.to_numpy() 114 | for i in range(-3, 8): 115 | check_is_close( 116 | nm.diagonal(arr, offset=i), 117 | np.diagonal(np_arr, offset=i), 118 | String("`diagonal` by axis {} is broken").format(i), 119 | ) 120 | -------------------------------------------------------------------------------- /tests/routines/test_manipulation.mojo: -------------------------------------------------------------------------------- 1 | import numojo as nm 2 | from numojo import * 3 | from testing.testing import assert_true, assert_almost_equal, assert_equal 4 | from utils_for_test import check, check_is_close 5 | from python import Python 6 | 7 | 8 | fn test_arr_manipulation() raises: 9 | var np = Python.import_module("numpy") 10 | 11 | # Test arange 12 | var A = nm.arange[nm.i16](1, 7, 1) 13 | var Anp = np.arange(1, 7, 1, dtype=np.int16) 14 | check_is_close(A, Anp, "Arange operation") 15 | 16 | var B = nm.random.randn(2, 3, 4) 17 | var Bnp = B.to_numpy() 18 | 19 | # Test flip 20 | check_is_close(nm.flip(B), np.flip(Bnp), "`flip` without `axis` fails.") 21 | for i in range(3): 22 | check_is_close( 23 | nm.flip(B, axis=i), 24 | np.flip(Bnp, axis=i), 25 | String("`flip` by `axis` {} fails.").format(i), 26 | ) 27 | 28 | 29 | def test_ravel_reshape(): 30 | var np = Python.import_module("numpy") 31 | var c = nm.fromstring[i8]( 32 | "[[[1,2,3,4][5,6,7,8]][[9,10,11,12][13,14,15,16]]]", order="C" 33 | ) 34 | var cnp = c.to_numpy() 35 | var f = nm.fromstring[i8]( 36 | "[[[1,2,3,4][5,6,7,8]][[9,10,11,12][13,14,15,16]]]", order="F" 37 | ) 38 | var fnp = f.to_numpy() 39 | 40 | # Test ravel 41 | check_is_close( 42 | nm.ravel(c, order="C"), 43 | np.ravel(cnp, order="C"), 44 | "`ravel` C-order array by C order is broken.", 45 | ) 46 | check_is_close( 47 | nm.ravel(c, order="F"), 48 | np.ravel(cnp, order="F"), 49 | "`ravel` C-order array by F order is broken.", 50 | ) 51 | check_is_close( 52 | nm.ravel(f, order="C"), 53 | np.ravel(fnp, order="C"), 54 | "`ravel` F-order array by C order is broken.", 55 | ) 56 | check_is_close( 57 | nm.ravel(f, order="F"), 58 | np.ravel(fnp, order="F"), 59 | "`ravel` F-order array by F order is broken.", 60 | ) 61 | 62 | # Test reshape 63 | check_is_close( 64 | nm.reshape(c, Shape(4, 2, 2), "C"), 65 | np.reshape(cnp, (4, 2, 2), "C"), 66 | "`reshape` C by C is broken", 67 | ) 68 | check_is_close( 69 | nm.reshape(c, Shape(4, 2, 2), "F"), 70 | np.reshape(cnp, (4, 2, 2), "F"), 71 | "`reshape` C by F is broken", 72 | ) 73 | check_is_close( 74 | nm.reshape(f, Shape(4, 2, 2), "C"), 75 | np.reshape(fnp, (4, 2, 2), "C"), 76 | "`reshape` F by C is broken", 77 | ) 78 | check_is_close( 79 | nm.reshape(f, Shape(4, 2, 2), "F"), 80 | np.reshape(fnp, (4, 2, 2), "F"), 81 | "`reshape` F by F is broken", 82 | ) 83 | 84 | 85 | def test_transpose(): 86 | var np = Python.import_module("numpy") 87 | var A = nm.random.randn(2) 88 | var Anp = A.to_numpy() 89 | check_is_close( 90 | nm.transpose(A), np.transpose(Anp), "1-d `transpose` is broken." 91 | ) 92 | A = nm.random.randn(2, 3) 93 | Anp = A.to_numpy() 94 | check_is_close( 95 | nm.transpose(A), np.transpose(Anp), "2-d `transpose` is broken." 96 | ) 97 | A = nm.random.randn(2, 3, 4) 98 | Anp = A.to_numpy() 99 | check_is_close( 100 | nm.transpose(A), np.transpose(Anp), "3-d `transpose` is broken." 101 | ) 102 | A = nm.random.randn(2, 3, 4, 5) 103 | Anp = A.to_numpy() 104 | check_is_close( 105 | nm.transpose(A), np.transpose(Anp), "4-d `transpose` is broken." 106 | ) 107 | check_is_close( 108 | A.T(), np.transpose(Anp), "4-d `transpose` with `.T` is broken." 109 | ) 110 | check_is_close( 111 | nm.transpose(A, axes=List(1, 3, 0, 2)), 112 | np.transpose(Anp, [1, 3, 0, 2]), 113 | "4-d `transpose` with arbitrary `axes` is broken.", 114 | ) 115 | 116 | 117 | def test_broadcast(): 118 | var np = Python.import_module("numpy") 119 | var a = nm.random.rand(Shape(2, 1, 3)) 120 | var Anp = a.to_numpy() 121 | check( 122 | nm.broadcast_to(a, Shape(2, 2, 3)), 123 | np.broadcast_to(a.to_numpy(), (2, 2, 3)), 124 | "`broadcast_to` fails.", 125 | ) 126 | check( 127 | nm.broadcast_to(a, Shape(2, 2, 2, 3)), 128 | np.broadcast_to(a.to_numpy(), (2, 2, 2, 3)), 129 | "`broadcast_to` fails.", 130 | ) 131 | -------------------------------------------------------------------------------- /tests/routines/test_random.mojo: -------------------------------------------------------------------------------- 1 | from math import sqrt 2 | import numojo as nm 3 | from numojo.prelude import * 4 | from python import Python, PythonObject 5 | from utils_for_test import check, check_is_close 6 | from testing.testing import assert_true, assert_almost_equal 7 | 8 | 9 | def test_rand(): 10 | """Test random array generation with specified shape.""" 11 | var arr = nm.random.rand[nm.f64](3, 5, 2) 12 | assert_true(arr.shape[0] == 3, "Shape of random array") 13 | assert_true(arr.shape[1] == 5, "Shape of random array") 14 | assert_true(arr.shape[2] == 2, "Shape of random array") 15 | 16 | 17 | def test_randminmax(): 18 | """Test random array generation with min and max values.""" 19 | var arr_variadic = nm.random.rand[nm.f64](10, 10, 10, min=1, max=2) 20 | var arr_list = nm.random.rand[nm.f64](List[Int](10, 10, 10), min=3, max=4) 21 | var arr_variadic_mean = nm.mean(arr_variadic) 22 | var arr_list_mean = nm.mean(arr_list) 23 | assert_almost_equal( 24 | arr_variadic_mean, 25 | 1.5, 26 | msg="Mean of random array within min and max", 27 | atol=0.1, 28 | ) 29 | assert_almost_equal( 30 | arr_list_mean, 31 | 3.5, 32 | msg="Mean of random array within min and max", 33 | atol=0.1, 34 | ) 35 | 36 | 37 | def test_randint(): 38 | """Test random int array generation with min and max values.""" 39 | var arr_low_high = nm.random.randint(Shape(10, 10, 10), 0, 10) 40 | var arr_high = nm.random.randint(Shape(10, 10, 10), 6) 41 | var arr_low_high_mean = nm.mean(arr_low_high) 42 | var arr_high_mean = nm.mean(arr_high) 43 | assert_almost_equal( 44 | arr_low_high_mean, 45 | 4.5, 46 | msg="Mean of `nm.random.randint(Shape(10, 10), 0, 10)` breaks", 47 | atol=0.1, 48 | ) 49 | assert_almost_equal( 50 | arr_high_mean, 51 | 2.5, 52 | msg="Mean of `nm.random.randint(Shape(10, 10), 6)` breaks", 53 | atol=0.1, 54 | ) 55 | 56 | 57 | def test_randn(): 58 | """Test random array generation with normal distribution.""" 59 | var arr_variadic_01 = nm.random.randn[nm.f64](20, 20, 20) 60 | var arr_variadic_31 = nm.random.randn[nm.f64]( 61 | Shape(20, 20, 20), mean=3, variance=1 62 | ) 63 | var arr_variadic_12 = nm.random.randn[nm.f64](Shape(20, 20, 20), 1, 2) 64 | 65 | var arr_variadic_mean01 = nm.mean(arr_variadic_01) 66 | var arr_variadic_mean31 = nm.mean(arr_variadic_31) 67 | var arr_variadic_mean12 = nm.mean(arr_variadic_12) 68 | var arr_variadic_var01 = nm.variance(arr_variadic_01) 69 | var arr_variadic_var31 = nm.variance(arr_variadic_31) 70 | var arr_variadic_var12 = nm.variance(arr_variadic_12) 71 | 72 | assert_almost_equal( 73 | arr_variadic_mean01, 74 | 0, 75 | msg="Mean of random array with mean 0 and variance 1", 76 | atol=0.1, 77 | ) 78 | assert_almost_equal( 79 | arr_variadic_mean31, 80 | 3, 81 | msg="Mean of random array with mean 3 and variance 1", 82 | atol=0.1, 83 | ) 84 | assert_almost_equal( 85 | arr_variadic_mean12, 86 | 1, 87 | msg="Mean of random array with mean 1 and variance 2", 88 | atol=0.1, 89 | ) 90 | 91 | assert_almost_equal( 92 | arr_variadic_var01, 93 | 1, 94 | msg="Variance of random array with mean 0 and variance 1", 95 | atol=0.1, 96 | ) 97 | assert_almost_equal( 98 | arr_variadic_var31, 99 | 1, 100 | msg="Variance of random array with mean 3 and variance 1", 101 | atol=0.1, 102 | ) 103 | assert_almost_equal( 104 | arr_variadic_var12, 105 | 2, 106 | msg="Variance of random array with mean 1 and variance 2", 107 | atol=0.1, 108 | ) 109 | 110 | 111 | def test_randn_list(): 112 | """Test random array generation with normal distribution.""" 113 | var arr_list_01 = nm.random.randn[nm.f64](Shape(20, 20, 20)) 114 | var arr_list_31 = nm.random.randn[nm.f64](Shape(20, 20, 20)) + 3 115 | var arr_list_12 = nm.random.randn[nm.f64](Shape(20, 20, 20)) * sqrt(2.0) + 1 116 | 117 | var arr_list_mean01 = nm.mean(arr_list_01) 118 | var arr_list_mean31 = nm.mean(arr_list_31) 119 | var arr_list_mean12 = nm.mean(arr_list_12) 120 | var arr_list_var01 = nm.variance(arr_list_01) 121 | var arr_list_var31 = nm.variance(arr_list_31) 122 | var arr_list_var12 = nm.variance(arr_list_12) 123 | 124 | assert_almost_equal( 125 | arr_list_mean01, 126 | 0, 127 | msg="Mean of random array with mean 0 and variance 1", 128 | atol=0.1, 129 | ) 130 | assert_almost_equal( 131 | arr_list_mean31, 132 | 3, 133 | msg="Mean of random array with mean 3 and variance 1", 134 | atol=0.1, 135 | ) 136 | assert_almost_equal( 137 | arr_list_mean12, 138 | 1, 139 | msg="Mean of random array with mean 1 and variance 2", 140 | atol=0.1, 141 | ) 142 | 143 | assert_almost_equal( 144 | arr_list_var01, 145 | 1, 146 | msg="Variance of random array with mean 0 and variance 1", 147 | atol=0.1, 148 | ) 149 | assert_almost_equal( 150 | arr_list_var31, 151 | 1, 152 | msg="Variance of random array with mean 3 and variance 1", 153 | atol=0.1, 154 | ) 155 | assert_almost_equal( 156 | arr_list_var12, 157 | 2, 158 | msg="Variance of random array with mean 1 and variance 2", 159 | atol=0.1, 160 | ) 161 | 162 | 163 | def test_rand_exponential(): 164 | """Test random array generation with exponential distribution.""" 165 | var arr_variadic = nm.random.exponential[nm.f64]( 166 | Shape(20, 20, 20), scale=2.0 167 | ) 168 | var arr_list = nm.random.exponential[nm.f64]( 169 | List[Int](20, 20, 20), scale=0.5 170 | ) 171 | 172 | var arr_variadic_mean = nm.mean(arr_variadic) 173 | var arr_list_mean = nm.mean(arr_list) 174 | 175 | # For exponential distribution, mean = 1 / rate 176 | assert_almost_equal( 177 | arr_variadic_mean, 178 | 1 / 2, 179 | msg="Mean of exponential distribution with rate 2.0", 180 | atol=0.1, 181 | ) 182 | assert_almost_equal( 183 | arr_list_mean, 184 | 1 / 0.5, 185 | msg="Mean of exponential distribution with rate 0.5", 186 | atol=0.2, 187 | ) 188 | 189 | # For exponential distribution, variance = 1 / (rate^2) 190 | var arr_variadic_var = nm.variance(arr_variadic) 191 | var arr_list_var = nm.variance(arr_list) 192 | 193 | assert_almost_equal( 194 | arr_variadic_var, 195 | 1 / 2**2, 196 | msg="Variance of exponential distribution with rate 2.0", 197 | atol=0.1, 198 | ) 199 | assert_almost_equal( 200 | arr_list_var, 201 | 1 / 0.5**2, 202 | msg="Variance of exponential distribution with rate 0.5", 203 | atol=0.5, 204 | ) 205 | 206 | # Test that all values are non-negative 207 | for i in range(arr_variadic.num_elements()): 208 | assert_true( 209 | arr_variadic._buf.ptr[i] >= 0, 210 | "Exponential distribution should only produce non-negative values", 211 | ) 212 | 213 | for i in range(arr_list.num_elements()): 214 | assert_true( 215 | arr_list._buf.ptr[i] >= 0, 216 | "Exponential distribution should only produce non-negative values", 217 | ) 218 | -------------------------------------------------------------------------------- /tests/routines/test_sorting.mojo: -------------------------------------------------------------------------------- 1 | import numojo as nm 2 | from python import Python, PythonObject 3 | from utils_for_test import check, check_is_close 4 | 5 | 6 | fn test_sorting() raises: 7 | var np = Python.import_module("numpy") 8 | var A = nm.random.randn(10) 9 | var B = nm.random.randn(2, 3) 10 | var C = nm.random.randn(2, 3, 4) 11 | var S = nm.random.randn(3, 3, 3, 3, 3, 3) # 6d array 12 | var Sf = S.reshape(S.shape, order="F") # 6d array F order 13 | 14 | # Sort 15 | check( 16 | nm.sort(A, axis=0), np.sort(A.to_numpy(), axis=0), "`sort` 1d is broken" 17 | ) 18 | check( 19 | nm.sort(B, axis=0), 20 | np.sort(B.to_numpy(), axis=0), 21 | "`sort` 2d by axis 0 is broken", 22 | ) 23 | check( 24 | nm.sort(B, axis=1), 25 | np.sort(B.to_numpy(), axis=1), 26 | "`sort` 2d by axis 1 is broken", 27 | ) 28 | for i in range(3): 29 | check( 30 | nm.sort(C, axis=i), 31 | np.sort(C.to_numpy(), axis=i), 32 | String("`sort` 3d by axis {} is broken").format(i), 33 | ) 34 | for i in range(6): 35 | check( 36 | nm.sort(S, axis=i), 37 | np.sort(S.to_numpy(), axis=i), 38 | String("`sort` 6d by axis {} is broken").format(i), 39 | ) 40 | 41 | # Argsort 42 | check( 43 | nm.argsort(A, axis=0), 44 | np.argsort(A.to_numpy(), axis=0), 45 | "`argsort` 1d is broken", 46 | ) 47 | check( 48 | nm.argsort(B, axis=0), 49 | np.argsort(B.to_numpy(), axis=0), 50 | "`argsort` 2d by axis 0 is broken", 51 | ) 52 | check( 53 | nm.argsort(B, axis=1), 54 | np.argsort(B.to_numpy(), axis=1), 55 | "`argsort` 2d by axis 1 is broken", 56 | ) 57 | for i in range(3): 58 | check_is_close( 59 | nm.argsort(C, axis=i), 60 | np.argsort(C.to_numpy(), axis=i), 61 | String("`argsort` 3d by axis {} is broken").format(i), 62 | ) 63 | check( 64 | nm.argsort(S), 65 | np.argsort(S.to_numpy(), axis=None), 66 | String("`argsort` 6d by axis None is broken"), 67 | ) 68 | for i in range(6): 69 | check( 70 | nm.argsort(S, axis=i), 71 | np.argsort(S.to_numpy(), axis=i), 72 | String("`argsort` 6d by axis {} is broken").format(i), 73 | ) 74 | 75 | check( 76 | nm.argsort(Sf), 77 | np.argsort(Sf.to_numpy(), axis=None), 78 | String("`argsort` 6d F-order array by axis None is broken"), 79 | ) 80 | for i in range(6): 81 | check( 82 | nm.argsort(Sf, axis=i), 83 | np.argsort(Sf.to_numpy(), axis=i), 84 | String("`argsort` 6d F-order by axis {} is broken").format(i), 85 | ) 86 | 87 | # In-place sort 88 | for i in range(3): 89 | C.sort(axis=i) 90 | Cnp = C.to_numpy() 91 | Cnp.sort(axis=i) 92 | check( 93 | C, 94 | Cnp, 95 | String("`NDArray.sort()` 3d by axis {} is broken").format(i), 96 | ) 97 | -------------------------------------------------------------------------------- /tests/routines/test_statistics.mojo: -------------------------------------------------------------------------------- 1 | import numojo as nm 2 | from numojo.prelude import * 3 | from numojo.core.matrix import Matrix 4 | from python import Python, PythonObject 5 | from testing.testing import assert_raises, assert_true 6 | from utils_for_test import check, check_is_close 7 | 8 | # ===-----------------------------------------------------------------------===# 9 | # Statistics 10 | # ===-----------------------------------------------------------------------===# 11 | 12 | 13 | def test_mean_median_var_std(): 14 | var np = Python.import_module("numpy") 15 | var sp = Python.import_module("scipy") 16 | var A = nm.random.randn(3, 4, 5) 17 | var Anp = A.to_numpy() 18 | 19 | assert_true( 20 | np.all(np.isclose(nm.mean(A), np.mean(Anp), atol=0.001)), 21 | "`mean` is broken", 22 | ) 23 | for axis in range(3): 24 | check_is_close( 25 | nm.mean(A, axis=axis), 26 | np.mean(Anp, axis=axis), 27 | String("`mean` is broken for axis {}".format(axis)), 28 | ) 29 | 30 | assert_true( 31 | np.all(np.isclose(nm.median(A), np.median(Anp), atol=0.001)), 32 | "`median` is broken", 33 | ) 34 | for axis in range(3): 35 | check_is_close( 36 | nm.median(A, axis), 37 | np.median(Anp, axis), 38 | String("`median` is broken for axis {}".format(axis)), 39 | ) 40 | 41 | assert_true( 42 | np.all( 43 | np.isclose( 44 | nm.mode(A), sp.stats.mode(Anp, axis=None).mode, atol=0.001 45 | ) 46 | ), 47 | "`mode` is broken", 48 | ) 49 | for axis in range(3): 50 | check_is_close( 51 | nm.mode(A, axis), 52 | sp.stats.mode(Anp, axis).mode, 53 | String("`mode` is broken for axis {}".format(axis)), 54 | ) 55 | 56 | assert_true( 57 | np.all(np.isclose(nm.variance(A), np.`var`(Anp), atol=0.001)), 58 | "`variance` is broken", 59 | ) 60 | for axis in range(3): 61 | check_is_close( 62 | nm.variance(A, axis), 63 | np.`var`(Anp, axis), 64 | String("`variance` is broken for axis {}".format(axis)), 65 | ) 66 | 67 | assert_true( 68 | np.all(np.isclose(nm.std(A), np.std(Anp), atol=0.001)), 69 | "`std` is broken", 70 | ) 71 | for axis in range(3): 72 | check_is_close( 73 | nm.std(A, axis), 74 | np.std(Anp, axis), 75 | String("`std` is broken for axis {}".format(axis)), 76 | ) 77 | -------------------------------------------------------------------------------- /tests/science/test_signal.mojo: -------------------------------------------------------------------------------- 1 | import numojo as nm 2 | from numojo.prelude import * 3 | from python import Python, PythonObject 4 | from utils_for_test import check, check_is_close 5 | from testing.testing import assert_raises 6 | 7 | 8 | def test_convolve2d(): 9 | var sp = Python.import_module("scipy") 10 | in1 = nm.random.rand(6, 6) 11 | in2 = nm.fromstring("[[1, 0], [0, -1]]") 12 | 13 | npin1 = in1.to_numpy() 14 | npin2 = in2.to_numpy() 15 | 16 | res1 = nm.science.signal.convolve2d(in1, in2) 17 | res2 = sp.signal.convolve2d(npin1, npin2, mode="valid") 18 | check(res1, res2, "test_convolve2d failed #2\n") 19 | -------------------------------------------------------------------------------- /tests/utils_for_test.mojo: -------------------------------------------------------------------------------- 1 | from python import Python, PythonObject 2 | from testing.testing import assert_true 3 | import numojo as nm 4 | 5 | 6 | fn check[ 7 | dtype: DType 8 | ](array: nm.NDArray[dtype], np_sol: PythonObject, st: String) raises: 9 | var np = Python.import_module("numpy") 10 | assert_true(np.all(np.equal(array.to_numpy(), np_sol)), st) 11 | 12 | 13 | fn check_with_dtype[ 14 | dtype: DType 15 | ](array: nm.NDArray[dtype], np_sol: PythonObject, st: String) raises: 16 | var np = Python.import_module("numpy") 17 | assert_true(String(array.dtype) == String(np_sol.dtype), "DType mismatch") 18 | assert_true(np.all(np.equal(array.to_numpy(), np_sol)), st) 19 | 20 | 21 | fn check_is_close[ 22 | dtype: DType 23 | ](array: nm.NDArray[dtype], np_sol: PythonObject, st: String) raises: 24 | var np = Python.import_module("numpy") 25 | assert_true(np.all(np.isclose(array.to_numpy(), np_sol, atol=0.1)), st) 26 | 27 | 28 | fn check_values_close[ 29 | dtype: DType 30 | ](value: Scalar[dtype], np_sol: PythonObject, st: String) raises: 31 | var np = Python.import_module("numpy") 32 | assert_true(np.isclose(value, np_sol, atol=0.001), st) 33 | --------------------------------------------------------------------------------