Mobius and mertens as symbolic functions

symengine · May 16, 2021 · ab5480c · ab5480c
1 parent 2ac1583
commit ab5480c
Show file tree

Hide file tree

Showing 13 changed files with 207 additions and 88 deletions.
diff --git a/benchmarks/ntheorybench.cpp b/benchmarks/ntheorybench.cpp
@@ -2,6 +2,7 @@
 #include <iostream>
 
 #include <symengine/ntheory.h>
+#include <symengine/ntheory_funcs.h>
 using std::cout;
 using std::endl;
 
@@ -16,7 +17,7 @@ void _bench_mertens(const unsigned long a)
 
     cout << "mertens(" << a << "):";
     t1 = std::chrono::high_resolution_clock::now();
-    mertens(a);
+    mertens(integer(a));
     t2 = std::chrono::high_resolution_clock::now();
     cout << std::chrono::duration_cast<std::chrono::milliseconds>(t2 - t1)
                 .count()
@@ -41,7 +42,7 @@ void _bench_mobius(const unsigned long a)
 {
     cout << "mobius(" << a << "): ";
     auto t1 = std::chrono::high_resolution_clock::now();
-    mobius(*integer(a));
+    mobius(integer(a));
     auto t2 = std::chrono::high_resolution_clock::now();
     cout << std::chrono::duration_cast<std::chrono::milliseconds>(t2 - t1)
                 .count()
@@ -102,4 +103,4 @@ int main()
     bench_mobius();
     bench_prime_factor_multiplicities();
     bench_mp_sqrt();
-}
+}
diff --git a/symengine/ntheory.cpp b/symengine/ntheory.cpp
@@ -1599,37 +1599,4 @@ i.e a % mod in set([i**n % mod for i in range(mod)]).
     return true;
 }
 
-int mobius(const Integer &a)
-{
-    if (a.as_int() <= 0) {
-        throw SymEngineException("mobius: Integer <= 0");
-    }
-    map_integer_uint prime_mul;
-    bool is_square_free = true;
-    prime_factor_multiplicities(prime_mul, a);
-    auto num_prime_factors = prime_mul.size();
-    for (const auto &it : prime_mul) {
-        int p_freq = it.second;
-        if (p_freq > 1) {
-            is_square_free = false;
-            break;
-        }
-    }
-    if (!is_square_free) {
-        return 0;
-    } else if (num_prime_factors % 2 == 0) {
-        return 1;
-    } else {
-        return -1;
-    }
-}
-
-long mertens(const unsigned long a)
-{
-    long mertens = 0;
-    for (unsigned long i = 1; i <= a; ++i) {
-        mertens += mobius(*(integer(i)));
-    }
-    return mertens;
-}
 } // namespace SymEngine
diff --git a/symengine/ntheory.h b/symengine/ntheory.h
@@ -151,15 +151,5 @@ vec_integer_class quadratic_residues(const Integer &a);
 bool is_quad_residue(const Integer &a, const Integer &p);
 //! Returns true if 'a' is a nth power residue of 'mod'
 bool is_nth_residue(const Integer &a, const Integer &n, const Integer &mod);
-//! Mobius Function
-// mu(n) = 1 if n is a square-free positive integer with an even number of prime
-// factors
-// mu(n) = −1 if n is a square-free positive integer with an odd number of prime
-// factors
-// mu(n) = 0 if n has a squared prime factor
-int mobius(const Integer &a);
-// Mertens Function
-// mertens(n) -> Sum of mobius(i) for i from 1 to n
-long mertens(const unsigned long a);
 } // namespace SymEngine
 #endif
diff --git a/symengine/ntheory_funcs.cpp b/symengine/ntheory_funcs.cpp
@@ -1,6 +1,7 @@
 #include <symengine/ntheory.h>
 #include <symengine/ntheory_funcs.h>
 #include <symengine/prime_sieve.h>
+#include <symengine/add.h>
 
 namespace SymEngine
 {
@@ -99,4 +100,99 @@ RCP<const Basic> primorial(const RCP<const Basic> &arg)
     return make_rcp<const Primorial>(arg);
 }
 
+Mobius::Mobius(const RCP<const Basic> &arg) : OneArgFunction(arg)
+{
+    SYMENGINE_ASSIGN_TYPEID()
+    SYMENGINE_ASSERT(is_canonical(arg))
+}
+
+bool Mobius::is_canonical(const RCP<const Basic> &arg) const
+{
+    if (is_a_Number(*arg) or is_a<Constant>(*arg)) {
+        return false;
+    } else {
+        return true;
+    }
+}
+
+RCP<const Basic> Mobius::create(const RCP<const Basic> &arg) const
+{
+    return mobius(arg);
+}
+
+RCP<const Basic> mobius(const RCP<const Basic> &arg)
+{
+    if (is_a_Number(*arg)) {
+        if (is_a<Integer>(*arg)
+            and down_cast<const Integer &>(*arg).is_positive()) {
+            map_integer_uint prime_mul;
+            bool is_square_free = true;
+            prime_factor_multiplicities(prime_mul,
+                                        down_cast<const Integer &>(*arg));
+            auto num_prime_factors = prime_mul.size();
+            for (const auto &it : prime_mul) {
+                int p_freq = it.second;
+                if (p_freq > 1) {
+                    is_square_free = false;
+                    break;
+                }
+            }
+            if (!is_square_free) {
+                return zero;
+            } else if (num_prime_factors % 2 == 0) {
+                return one;
+            } else {
+                return minus_one;
+            }
+        } else if (is_a<NaN>(*arg)) {
+            return arg;
+        } else {
+            throw SymEngineException(
+                "Only positive integers are allowed by the Möbius function!");
+        }
+    }
+    return make_rcp<const Mobius>(arg);
+}
+
+Mertens::Mertens(const RCP<const Basic> &arg) : OneArgFunction(arg)
+{
+    SYMENGINE_ASSIGN_TYPEID()
+    SYMENGINE_ASSERT(is_canonical(arg))
+}
+
+bool Mertens::is_canonical(const RCP<const Basic> &arg) const
+{
+    if (is_a_Number(*arg) or is_a<Constant>(*arg)) {
+        return false;
+    } else {
+        return true;
+    }
+}
+
+RCP<const Basic> Mertens::create(const RCP<const Basic> &arg) const
+{
+    return mertens(arg);
+}
+
+RCP<const Basic> mertens(const RCP<const Basic> &arg)
+{
+    if (is_a_Number(*arg)) {
+        if (is_a<Integer>(*arg)
+            and down_cast<const Integer &>(*arg).is_positive()) {
+            RCP<const Basic> s = make_rcp<const Integer>(0);
+            auto a = down_cast<const Integer &>(*arg).as_uint();
+            for (unsigned long i = 1; i <= a; ++i) {
+                s = add(s, mobius(integer(i)));
+            }
+            return s;
+        } else if (is_a<NaN>(*arg)) {
+            return arg;
+        } else {
+            throw SymEngineException(
+                "Only positive integers are allowed by the Mertens function!");
+        }
+    }
+    return make_rcp<const Mertens>(arg);
+}
+
 } // namespace SymEngine
diff --git a/symengine/ntheory_funcs.h b/symengine/ntheory_funcs.h
@@ -42,6 +42,36 @@ class Primorial : public OneArgFunction
 
 RCP<const Basic> primorial(const RCP<const Basic> &arg);
 
+class Mobius : public OneArgFunction
+{
+    /*! The Möbius mu function
+     * https://en.wikipedia.org/wiki/Möbius_function
+     **/
+
+public:
+    IMPLEMENT_TYPEID(SYMENGINE_MOBIUS)
+    Mobius(const RCP<const Basic> &arg);
+    bool is_canonical(const RCP<const Basic> &arg) const;
+    RCP<const Basic> create(const RCP<const Basic> &arg) const;
+};
+
+RCP<const Basic> mobius(const RCP<const Basic> &arg);
+
+class Mertens : public OneArgFunction
+{
+    /*! The Mertens function (summatory moebius)
+     * https://en.wikipedia.org/wiki/Mertens_function
+     **/
+
+public:
+    IMPLEMENT_TYPEID(SYMENGINE_MERTENS)
+    Mertens(const RCP<const Basic> &arg);
+    bool is_canonical(const RCP<const Basic> &arg) const;
+    RCP<const Basic> create(const RCP<const Basic> &arg) const;
+};
+
+RCP<const Basic> mertens(const RCP<const Basic> &arg);
+
 } // namespace SymEngine
 
 #endif
diff --git a/symengine/parser/parser.cpp b/symengine/parser/parser.cpp
@@ -100,6 +100,8 @@ RCP<const Basic> Parser::functionify(const std::string &name, vec_basic &params)
             {"zeta", single_casted_zeta},
             {"primepi", primepi},
             {"primorial", primorial},
+            {"mobius", mobius},
+            {"mertens", mertens},
         };
     const static std::map<
         const std::string,

diff --git a/symengine/printers/latex.cpp b/symengine/printers/latex.cpp
@@ -443,6 +443,8 @@ std::vector<std::string> init_latex_printer_names()
     names[SYMENGINE_GAMMA] = "\\Gamma";
     names[SYMENGINE_TRUNCATE] = "\\operatorname{truncate}";
     names[SYMENGINE_PRIMEPI] = "\\pi";
+    names[SYMENGINE_MOBIUS] = "\\mu";
+    names[SYMENGINE_MERTENS] = "\\operatorname{M}";
     return names;
 }
 

diff --git a/symengine/printers/strprinter.cpp b/symengine/printers/strprinter.cpp
@@ -1089,6 +1089,8 @@ std::vector<std::string> init_str_printer_names()
     names[SYMENGINE_CONJUGATE] = "conjugate";
     names[SYMENGINE_PRIMEPI] = "primepi";
     names[SYMENGINE_PRIMORIAL] = "primorial";
+    names[SYMENGINE_MOBIUS] = "mobius";
+    names[SYMENGINE_MERTENS] = "mertens";
     names[SYMENGINE_UNEVALUATED_EXPR] = "";
     return names;
 }

diff --git a/symengine/tests/basic/test_parser.cpp b/symengine/tests/basic/test_parser.cpp
@@ -761,6 +761,14 @@ TEST_CASE("Parsing: function_symbols", "[parser]")
     s = "primorial(15.9)";
     res = parse(s);
     REQUIRE(eq(*res, *integer(30030)));
+
+    s = "mobius(28)";
+    res = parse(s);
+    REQUIRE(eq(*res, *integer(0)));
+
+    s = "mertens(28)";
+    res = parse(s);
+    REQUIRE(eq(*res, *integer(-1)));
 }
 
 TEST_CASE("Parsing: multi-arg functions", "[parser]")

diff --git a/symengine/tests/ntheory/test_ntheory.cpp b/symengine/tests/ntheory/test_ntheory.cpp
@@ -21,7 +21,6 @@ using SymEngine::integer_class;
 using SymEngine::is_a;
 using SymEngine::lucas;
 using SymEngine::map_integer_uint;
-using SymEngine::mertens;
 using SymEngine::minus_one;
 using SymEngine::mod;
 using SymEngine::mod_f;
@@ -844,47 +843,6 @@ TEST_CASE("test_is_nth_residue(): ntheory", "[ntheory]")
     REQUIRE(is_nth_residue(*i32, *i10, *im1) == true);
 }
 
-TEST_CASE("test_mobius(): ntheory", "[ntheory]")
-{
-    RCP<const Integer> i1 = integer(1);
-    RCP<const Integer> i2 = integer(2);
-    RCP<const Integer> i3 = integer(3);
-    RCP<const Integer> i4 = integer(4);
-    RCP<const Integer> i5 = integer(5);
-    RCP<const Integer> i6 = integer(6);
-    RCP<const Integer> i7 = integer(7);
-    RCP<const Integer> i8 = integer(8);
-    RCP<const Integer> i9 = integer(9);
-    RCP<const Integer> i10 = integer(10);
-
-    REQUIRE(mobius(*i1) == 1);
-    REQUIRE(mobius(*i2) == -1);
-    REQUIRE(mobius(*i3) == -1);
-    REQUIRE(mobius(*i4) == 0);
-    REQUIRE(mobius(*i5) == -1);
-    REQUIRE(mobius(*i6) == 1);
-    REQUIRE(mobius(*i7) == -1);
-    REQUIRE(mobius(*i8) == 0);
-    REQUIRE(mobius(*i9) == 0);
-    REQUIRE(mobius(*i10) == 1);
-
-    CHECK_THROWS_AS(mobius(*minus_one), SymEngineException &);
-}
-
-TEST_CASE("test_mertens(): ntheory", "[ntheory]")
-{
-    REQUIRE(mertens(1) == 1);
-    REQUIRE(mertens(2) == 0);
-    REQUIRE(mertens(4) == -1);
-    REQUIRE(mertens(12) == -2);
-    REQUIRE(mertens(13) == -3);
-    REQUIRE(mertens(22) == -1);
-    REQUIRE(mertens(31) == -4);
-    REQUIRE(mertens(36) == -1);
-    REQUIRE(mertens(39) == 0);
-    REQUIRE(mertens(113) == -5);
-}
-
 TEST_CASE("test_harmonic(): ntheory", "[ntheory]")
 {
     RCP<const Number> r1, r2;

diff --git a/symengine/tests/ntheory/test_ntheory_funcs.cpp b/symengine/tests/ntheory/test_ntheory_funcs.cpp
@@ -11,8 +11,14 @@ using SymEngine::down_cast;
 using SymEngine::Inf;
 using SymEngine::Infty;
 using SymEngine::integer;
+using SymEngine::Integer;
 using SymEngine::is_a;
 using SymEngine::make_rcp;
+using SymEngine::mertens;
+using SymEngine::Mertens;
+using SymEngine::minus_one;
+using SymEngine::mobius;
+using SymEngine::Mobius;
 using SymEngine::Nan;
 using SymEngine::NaN;
 using SymEngine::pi;
@@ -68,3 +74,54 @@ TEST_CASE("test_primorial(): ntheory_funcs", "[ntheory_funcs]")
     REQUIRE(eq(*Primorial(symbol("x")).create(integer(1)), *integer(1)));
     REQUIRE(not(Primorial(symbol("x")).is_canonical(integer(1))));
 }
+
+TEST_CASE("test_mobius(): ntheory", "[ntheory]")
+{
+    RCP<const Integer> i1 = integer(1);
+    RCP<const Integer> i2 = integer(2);
+    RCP<const Integer> i3 = integer(3);
+    RCP<const Integer> i4 = integer(4);
+    RCP<const Integer> i5 = integer(5);
+    RCP<const Integer> i6 = integer(6);
+    RCP<const Integer> i7 = integer(7);
+    RCP<const Integer> i8 = integer(8);
+    RCP<const Integer> i9 = integer(9);
+    RCP<const Integer> i10 = integer(10);
+
+    REQUIRE(eq(*mobius(i1), *integer(1)));
+    REQUIRE(eq(*mobius(i2), *integer(-1)));
+    REQUIRE(eq(*mobius(i3), *integer(-1)));
+    REQUIRE(eq(*mobius(i4), *integer(0)));
+    REQUIRE(eq(*mobius(i5), *integer(-1)));
+    REQUIRE(eq(*mobius(i6), *integer(1)));
+    REQUIRE(eq(*mobius(i7), *integer(-1)));
+    REQUIRE(eq(*mobius(i8), *integer(0)));
+    REQUIRE(eq(*mobius(i9), *integer(0)));
+    REQUIRE(eq(*mobius(i10), *integer(1)));
+
+    CHECK_THROWS_AS(mobius(minus_one), SymEngineException &);
+    REQUIRE(eq(*mobius(symbol("x")), *make_rcp<Mobius>(symbol("x"))));
+    REQUIRE(mobius(symbol("x"))->__str__() == "mobius(x)");
+    REQUIRE(eq(*Mobius(symbol("x")).create(integer(2)), *integer(-1)));
+    REQUIRE(not(Mobius(symbol("x")).is_canonical(integer(1))));
+}
+
+TEST_CASE("test_mertens(): ntheory", "[ntheory]")
+{
+    REQUIRE(eq(*mertens(integer(1)), *integer(1)));
+    REQUIRE(eq(*mertens(integer(2)), *integer(0)));
+    REQUIRE(eq(*mertens(integer(4)), *integer(-1)));
+    REQUIRE(eq(*mertens(integer(12)), *integer(-2)));
+    REQUIRE(eq(*mertens(integer(13)), *integer(-3)));
+    REQUIRE(eq(*mertens(integer(22)), *integer(-1)));
+    REQUIRE(eq(*mertens(integer(31)), *integer(-4)));
+    REQUIRE(eq(*mertens(integer(36)), *integer(-1)));
+    REQUIRE(eq(*mertens(integer(39)), *integer(0)));
+    REQUIRE(eq(*mertens(integer(113)), *integer(-5)));
+
+    CHECK_THROWS_AS(mertens(minus_one), SymEngineException &);
+    REQUIRE(eq(*mertens(symbol("x")), *make_rcp<Mertens>(symbol("x"))));
+    REQUIRE(mertens(symbol("x"))->__str__() == "mertens(x)");
+    REQUIRE(eq(*Mertens(symbol("x")).create(integer(2)), *integer(0)));
+    REQUIRE(not(Mertens(symbol("x")).is_canonical(integer(1))));
+}