colormap and possible normalize

matthew-hennefarth · matthew-hennefarth · commit b8a3f1060ee0 · 2021-05-14T22:36:30.000-07:00
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -128,4 +128,4 @@ if (BUILD_INSTALLER AND BUILD_PACKAGE)
 
     # Always include CPack at last
     include(CPack)
-endif ()
+endif ()
diff --git a/source/matplot/axes_objects/vectors.cpp b/source/matplot/axes_objects/vectors.cpp
@@ -48,12 +48,30 @@ namespace matplot {
           x_data_(x_data), z_data_(z_data), u_data_(u_data), v_data_(v_data),
           w_data_(w_data) {}
 
+    vectors::vectors(class axes_type *parent, const std::vector<double> &x_data,
+                     const std::vector<double> &y_data,
+                     const std::vector<double> &z_data,
+                     const std::vector<double> &u_data,
+                     const std::vector<double> &v_data,
+                     const std::vector<double> &w_data,
+                     const std::vector<double> &m_data,
+                     std::string_view line_spec)
+        : axes_object(parent), line_spec_(this, line_spec), y_data_(y_data),
+          x_data_(x_data), z_data_(z_data), u_data_(u_data), v_data_(v_data),
+          w_data_(w_data), m_data_(m_data) {}
+
     std::string vectors::plot_string() {
         maybe_update_line_spec();
         std::stringstream ss;
-        ss << " '-' with vectors " +
-                  line_spec_.plot_string(
-                      line_spec::style_to_plot::plot_line_only, false);
+        ss << " '-' with vectors";
+        if (!line_spec_.user_color()) {
+            ss << " linecolor palette ";
+        } else {
+            ss << " linecolor rgb \'" << to_string(line_spec_.color()) << "\'";
+        }
+
+        ss << " linewidth " << line_spec_.line_width();
+
         if (use_y2_) {
             ss << " xlim x1y2";
         }
@@ -79,6 +97,7 @@ namespace matplot {
                 double u_value = u_data_.size() > i ? u_data_[i] : 0;
                 double v_value = v_data_.size() > i ? v_data_[i] : 0;
                 double w_value = w_data_.size() > i ? w_data_[i] : 0;
+                double m_value = m_data_.size() > i ? m_data_[i] : 0;
 
                 bool is_end_of_series =
                     !std::isfinite(x_value) || !std::isfinite(y_value) ||
@@ -97,6 +116,15 @@ namespace matplot {
                     ss << "  " << 0.0;
                 }
 
+                if (normalize_ && m_value != 0) {
+                    double length = sqrt(u_value * u_value + v_value * v_value +
+                                         w_value * w_value);
+                    double scale = length / m_value;
+                    u_value *= scale / length;
+                    v_value *= scale / length;
+                    w_value *= scale / length;
+                }
+
                 ss << "  " << u_value;
                 ss << "  " << v_value;
                 if (is_3d()) {
@@ -105,6 +133,8 @@ namespace matplot {
                     ss << "  " << 0.0;
                 }
 
+                ss << " " << m_value;
+
                 ss << "\n";
             }
             // if polar, gnuplot uses the vector origin to determine
@@ -123,13 +153,8 @@ namespace matplot {
     }
 
     void vectors::maybe_update_line_spec() {
-        if (line_spec_.has_line() && !line_spec_.user_color()) {
-            // if user didn't set the color, get color from xlim
-            auto c = parent_->get_color_and_bump();
-            line_spec_.color(c);
-        } else if (line_spec_.has_non_custom_marker() &&
-                   !line_spec_.user_color() &&
-                   !line_spec_.marker_user_color()) {
+        if (line_spec_.has_non_custom_marker() && !line_spec_.user_color() &&
+            !line_spec_.marker_user_color()) {
             auto c = parent_->get_color_and_bump();
             line_spec_.marker_color(c);
         }
@@ -399,4 +424,12 @@ namespace matplot {
         return *this;
     }
 
-} // namespace matplot
+    bool vectors::norm() const { return normalize_; }
+
+    class vectors &vectors::norm(bool normalize) {
+        normalize_ = normalize;
+        touch();
+        return *this;
+    }
+
+} // namespace matplot
diff --git a/source/matplot/axes_objects/vectors.h b/source/matplot/axes_objects/vectors.h
@@ -48,6 +48,15 @@ namespace matplot {
                 const std::vector<double> &w_data,
                 std::string_view line_spec = "");
 
+        vectors(class axes_type *parent, const std::vector<double> &x_data,
+            const std::vector<double> &y_data,
+            const std::vector<double> &z_data,
+            const std::vector<double> &u_data,
+            const std::vector<double> &v_data,
+            const std::vector<double> &w_data,
+            const std::vector<double> &m_data,
+            std::string_view line_spec = "");
+
         /// If we receive an axes_handle, we can convert it to a raw
         /// pointer because there is no ownership involved here
         template <class... Args>
@@ -102,6 +111,9 @@ namespace matplot {
 
         bool visible() const;
         class vectors &visible(bool visible);
+        
+        bool norm() const;
+        class vectors &norm(bool normalize);
 
       public /* getters and setters bypassing the line_spec */:
         float line_width() const;
@@ -194,6 +206,9 @@ namespace matplot {
         std::vector<double> v_data_{};
         std::vector<double> w_data_{};
 
+        /// Vector magnitude
+        std::vector<double> m_data_{};
+
         /// Positions at which we want markers to appear
         std::vector<size_t> marker_indices_{};
         std::vector<float> marker_sizes_{};
@@ -213,7 +228,10 @@ namespace matplot {
 
         /// True if visible
         bool visible_{true};
-    };
+
+        /// True if vectors are normalized
+        bool normalize_{false};
+};
 } // namespace matplot
 
 #endif // MATPLOTPLUSPLUS_VECTORS_H
diff --git a/source/matplot/core/axes_type.cpp b/source/matplot/core/axes_type.cpp
